Merge pull request #309 from fireyang/master

fix rpc client call WARNING: DATA RACE

.travis.yml

@@ -1,7 +1,7 @@
 language: go
 go:
-- 1.9.5
 - 1.10.x
+- 1.11.x
 notifications:
   slack:
     secure: aEvhLbhujaGaKSrOokiG3//PaVHTIrc3fBpoRbCRqfZpyq6WREoapJJhF+tIpWWOwaC9GmChbD6aHo/jMUgwKXVyPSaNjiEL87YzUUpL8B2zslNp1rgfTg/LrzthOx3Q1TYwpaAl3to0fuHUVFX4yMeC2vuThq7WSXgMMxFCtbc=

README.md

@@ -192,7 +192,7 @@ func main() {
 	service.Init()
 
 	// Create new greeter client
-	greeter := proto.GreeterServiceClient("greeter", service.Client())
+	greeter := proto.NewGreeterService("greeter", service.Client())
 
 	// Call the greeter
 	rsp, err := greeter.Hello(context.TODO(), &proto.HelloRequest{Name: "John"})

broker/http_broker.go

@@ -276,12 +276,15 @@ func (h *httpBroker) Address() string {
 }
 
 func (h *httpBroker) Connect() error {
-	h.Lock()
-	defer h.Unlock()
-
+	h.RLock()
 	if h.running {
+		h.RUnlock()
 		return nil
 	}
+	h.RUnlock()
+
+	h.Lock()
+	defer h.Unlock()
 
 	var l net.Listener
 	var err error
@@ -326,10 +329,16 @@ func (h *httpBroker) Connect() error {
 	}
 
 	log.Logf("Broker Listening on %s", l.Addr().String())
+	addr := h.address
 	h.address = l.Addr().String()
 
 	go http.Serve(l, h.mux)
-	go h.run(l)
+	go func() {
+		h.run(l)
+		h.Lock()
+		h.address = addr
+		h.Unlock()
+	}()
 
 	// get registry
 	reg, ok := h.opts.Context.Value(registryKey).(registry.Registry)
@@ -345,12 +354,16 @@ func (h *httpBroker) Connect() error {
 }
 
 func (h *httpBroker) Disconnect() error {
-	h.Lock()
-	defer h.Unlock()
 
+	h.RLock()
 	if !h.running {
+		h.RUnlock()
 		return nil
 	}
+	h.RUnlock()
+
+	h.Lock()
+	defer h.Unlock()
 
 	// stop rcache
 	rc, ok := h.r.(rcache.Cache)
@@ -369,17 +382,24 @@ func (h *httpBroker) Disconnect() error {
 }
 
 func (h *httpBroker) Init(opts ...Option) error {
+	h.RLock()
+	if h.running {
+		h.RUnlock()
+		return errors.New("cannot init while connected")
+	}
+	h.RUnlock()
+
 	h.Lock()
 	defer h.Unlock()
 
-	if h.running {
-		return errors.New("cannot init while connected")
-	}
-
 	for _, o := range opts {
 		o(&h.opts)
 	}
 
+	if len(h.opts.Addrs) > 0 && len(h.opts.Addrs[0]) > 0 {
+		h.address = h.opts.Addrs[0]
+	}
+
 	if len(h.id) == 0 {
 		h.id = "broker-" + uuid.NewUUID().String()
 	}

client/client.go

@@ -12,7 +12,7 @@ import (
 type Client interface {
 	Init(...Option) error
 	Options() Options
-	NewMessage(topic string, msg interface{}) Message
+	NewMessage(topic string, msg interface{}, opts ...MessageOption) Message
 	NewRequest(service, method string, req interface{}, reqOpts ...RequestOption) Request
 	Call(ctx context.Context, req Request, rsp interface{}, opts ...CallOption) error
 	Stream(ctx context.Context, req Request, opts ...CallOption) (Stream, error)
@@ -56,6 +56,9 @@ type CallOption func(*CallOptions)
 // PublishOption used by Publish
 type PublishOption func(*PublishOptions)
 
+// MessageOption used by NewMessage
+type MessageOption func(*MessageOptions)
+
 // RequestOption used by NewRequest
 type RequestOption func(*RequestOptions)
 
@@ -65,13 +68,13 @@ var (
 	// DefaultBackoff is the default backoff function for retries
 	DefaultBackoff = exponentialBackoff
 	// DefaultRetry is the default check-for-retry function for retries
-	DefaultRetry = alwaysRetry
+	DefaultRetry = RetryOnError
 	// DefaultRetries is the default number of times a request is tried
 	DefaultRetries = 1
 	// DefaultRequestTimeout is the default request timeout
 	DefaultRequestTimeout = time.Second * 5
 	// DefaultPoolSize sets the connection pool size
-	DefaultPoolSize = 0
+	DefaultPoolSize = 1
 	// DefaultPoolTTL sets the connection pool ttl
 	DefaultPoolTTL = time.Minute
 )
@@ -83,13 +86,13 @@ func Call(ctx context.Context, request Request, response interface{}, opts ...Ca
 
 // Publishes a publication using the default client. Using the underlying broker
 // set within the options.
-func Publish(ctx context.Context, msg Message) error {
-	return DefaultClient.Publish(ctx, msg)
+func Publish(ctx context.Context, msg Message, opts ...PublishOption) error {
+	return DefaultClient.Publish(ctx, msg, opts...)
 }
 
 // Creates a new message using the default client
-func NewMessage(topic string, payload interface{}) Message {
-	return DefaultClient.NewMessage(topic, payload)
+func NewMessage(topic string, payload interface{}, opts ...MessageOption) Message {
+	return DefaultClient.NewMessage(topic, payload, opts...)
 }
 
 // Creates a new client with the options passed in

client/mock/mock.go

@@ -49,8 +49,8 @@ func (m *MockClient) Options() client.Options {
 	return m.Opts
 }
 
-func (m *MockClient) NewMessage(topic string, msg interface{}) client.Message {
-	return m.Client.NewMessage(topic, msg)
+func (m *MockClient) NewMessage(topic string, msg interface{}, opts ...client.MessageOption) client.Message {
+	return m.Client.NewMessage(topic, msg, opts...)
 }
 
 func (m *MockClient) NewRequest(service, method string, req interface{}, reqOpts ...client.RequestOption) client.Request {
@@ -80,8 +80,12 @@ func (m *MockClient) Call(ctx context.Context, req client.Request, rsp interface
 		if t := reflect.TypeOf(rsp); t.Kind() == reflect.Ptr {
 			v = reflect.Indirect(v)
 		}
+		response := r.Response
+		if t := reflect.TypeOf(r.Response); t.Kind() == reflect.Func {
+			response = reflect.ValueOf(r.Response).Call([]reflect.Value{})[0].Interface()
+		}
 
-		v.Set(reflect.ValueOf(r.Response))
+		v.Set(reflect.ValueOf(response))
 
 		return nil
 	}

client/mock/mock_test.go

@@ -16,6 +16,8 @@ func TestClient(t *testing.T) {
 		{Method: "Foo.Bar", Response: map[string]interface{}{"foo": "bar"}},
 		{Method: "Foo.Struct", Response: &TestResponse{Param: "aparam"}},
 		{Method: "Foo.Fail", Error: errors.InternalServerError("go.mock", "failed")},
+		{Method: "Foo.Func", Response: func() string { return "string" }},
+		{Method: "Foo.FuncStruct", Response: func() *TestResponse { return &TestResponse{Param: "aparam"} }},
 	}
 
 	c := NewClient(Response("go.mock", response))

client/options.go

@@ -67,6 +67,10 @@ type PublishOptions struct {
 	Context context.Context
 }
 
+type MessageOptions struct {
+	ContentType string
+}
+
 type RequestOptions struct {
 	ContentType string
 	Stream      bool

client/retry.go

@@ -2,12 +2,34 @@ package client
 
 import (
 	"context"
+
+	"github.com/micro/go-micro/errors"
 )
 
 // note that returning either false or a non-nil error will result in the call not being retried
 type RetryFunc func(ctx context.Context, req Request, retryCount int, err error) (bool, error)
 
-// always retry on error
-func alwaysRetry(ctx context.Context, req Request, retryCount int, err error) (bool, error) {
+// RetryAlways always retry on error
+func RetryAlways(ctx context.Context, req Request, retryCount int, err error) (bool, error) {
 	return true, nil
 }
+
+// RetryOnError retries a request on a 500 or timeout error
+func RetryOnError(ctx context.Context, req Request, retryCount int, err error) (bool, error) {
+	if err == nil {
+		return false, nil
+	}
+
+	e := errors.Parse(err.Error())
+	if e == nil {
+		return false, nil
+	}
+
+	switch e.Code {
+	// retry on timeout or internal server error
+	case 408, 500:
+		return true, nil
+	default:
+		return false, nil
+	}
+}

client/rpc_client.go

@@ -7,6 +7,8 @@ import (
 	"sync"
 	"time"
 
+	"sync/atomic"
+
 	"github.com/micro/go-micro/broker"
 	"github.com/micro/go-micro/codec"
 	"github.com/micro/go-micro/errors"
@@ -14,7 +16,6 @@ import (
 	"github.com/micro/go-micro/registry"
 	"github.com/micro/go-micro/selector"
 	"github.com/micro/go-micro/transport"
-	"sync/atomic"
 )
 
 type rpcClient struct {
@@ -88,7 +89,7 @@ func (r *rpcClient) call(ctx context.Context, address string, req Request, resp
 		r.pool.release(address, c, grr)
 	}()
 
-	seq := r.seq
+	seq := atomic.LoadUint64(&r.seq)
 	atomic.AddUint64(&r.seq, 1)
 
 	stream := &rpcStream{
@@ -159,7 +160,15 @@ func (r *rpcClient) stream(ctx context.Context, address string, req Request, opt
 		return nil, errors.InternalServerError("go.micro.client", err.Error())
 	}
 
-	c, err := r.opts.Transport.Dial(address, transport.WithStream(), transport.WithTimeout(opts.DialTimeout))
+	dOpts := []transport.DialOption{
+		transport.WithStream(),
+	}
+
+	if opts.DialTimeout >= 0 {
+		dOpts = append(dOpts, transport.WithTimeout(opts.DialTimeout))
+	}
+
+	c, err := r.opts.Transport.Dial(address, dOpts...)
 	if err != nil {
 		return nil, errors.InternalServerError("go.micro.client", "connection error: %v", err)
 	}
@@ -202,9 +211,12 @@ func (r *rpcClient) Init(opts ...Option) error {
 		o(&r.opts)
 	}
 
-	// recreate the pool if the options changed
+	// update pool configuration if the options changed
 	if size != r.opts.PoolSize || ttl != r.opts.PoolTTL {
-		r.pool = newPool(r.opts.PoolSize, r.opts.PoolTTL)
+		r.pool.Lock()
+		r.pool.size = r.opts.PoolSize
+		r.pool.ttl = int64(r.opts.PoolTTL.Seconds())
+		r.pool.Unlock()
 	}
 
 	return nil
@@ -227,9 +239,9 @@ func (r *rpcClient) next(request Request, opts CallOptions) (selector.Next, erro
 	// get next nodes from the selector
 	next, err := r.opts.Selector.Select(request.Service(), opts.SelectOptions...)
 	if err != nil && err == selector.ErrNotFound {
-		return nil, errors.NotFound("go.micro.client", err.Error())
+		return nil, errors.NotFound("go.micro.client", "service %s: %v", request.Service(), err.Error())
 	} else if err != nil {
-		return nil, errors.InternalServerError("go.micro.client", err.Error())
+		return nil, errors.InternalServerError("go.micro.client", "error selecting %s node: %v", request.Service(), err.Error())
 	}
 
 	return next, nil
@@ -279,7 +291,7 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 		// call backoff first. Someone may want an initial start delay
 		t, err := callOpts.Backoff(ctx, request, i)
 		if err != nil {
-			return errors.InternalServerError("go.micro.client", err.Error())
+			return errors.InternalServerError("go.micro.client", "backoff error: %v", err.Error())
 		}
 
 		// only sleep if greater than 0
@@ -290,9 +302,9 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 		// select next node
 		node, err := next()
 		if err != nil && err == selector.ErrNotFound {
-			return errors.NotFound("go.micro.client", err.Error())
+			return errors.NotFound("go.micro.client", "service %s: %v", request.Service(), err.Error())
 		} else if err != nil {
-			return errors.InternalServerError("go.micro.client", err.Error())
+			return errors.InternalServerError("go.micro.client", "error getting next %s node: %v", request.Service(), err.Error())
 		}
 
 		// set the address
@@ -310,10 +322,10 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 	ch := make(chan error, callOpts.Retries)
 	var gerr error
 
-	for i := 0; i < callOpts.Retries; i++ {
-		go func() {
+	for i := 0; i <= callOpts.Retries; i++ {
+		go func(i int) {
 			ch <- call(i)
-		}()
+		}(i)
 
 		select {
 		case <-ctx.Done():
@@ -352,18 +364,6 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, opts ...CallOpt
 		return nil, err
 	}
 
-	// check if we already have a deadline
-	d, ok := ctx.Deadline()
-	if !ok {
-		// no deadline so we create a new one
-		ctx, _ = context.WithTimeout(ctx, callOpts.RequestTimeout)
-	} else {
-		// got a deadline so no need to setup context
-		// but we need to set the timeout we pass along
-		opt := WithRequestTimeout(d.Sub(time.Now()))
-		opt(&callOpts)
-	}
-
 	// should we noop right here?
 	select {
 	case <-ctx.Done():
@@ -375,7 +375,7 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, opts ...CallOpt
 		// call backoff first. Someone may want an initial start delay
 		t, err := callOpts.Backoff(ctx, request, i)
 		if err != nil {
-			return nil, errors.InternalServerError("go.micro.client", err.Error())
+			return nil, errors.InternalServerError("go.micro.client", "backoff error: %v", err.Error())
 		}
 
 		// only sleep if greater than 0
@@ -385,9 +385,9 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, opts ...CallOpt
 
 		node, err := next()
 		if err != nil && err == selector.ErrNotFound {
-			return nil, errors.NotFound("go.micro.client", err.Error())
+			return nil, errors.NotFound("go.micro.client", "service %s: %v", request.Service(), err.Error())
 		} else if err != nil {
-			return nil, errors.InternalServerError("go.micro.client", err.Error())
+			return nil, errors.InternalServerError("go.micro.client", "error getting next %s node: %v", request.Service(), err.Error())
 		}
 
 		address := node.Address
@@ -408,7 +408,7 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, opts ...CallOpt
 	ch := make(chan response, callOpts.Retries)
 	var grr error
 
-	for i := 0; i < callOpts.Retries; i++ {
+	for i := 0; i <= callOpts.Retries; i++ {
 		go func() {
 			s, err := call(i)
 			ch <- response{s, err}
@@ -465,8 +465,8 @@ func (r *rpcClient) Publish(ctx context.Context, msg Message, opts ...PublishOpt
 	})
 }
 
-func (r *rpcClient) NewMessage(topic string, message interface{}) Message {
-	return newMessage(topic, message, r.opts.ContentType)
+func (r *rpcClient) NewMessage(topic string, message interface{}, opts ...MessageOption) Message {
+	return newMessage(topic, message, r.opts.ContentType, opts...)
 }
 
 func (r *rpcClient) NewRequest(service, method string, request interface{}, reqOpts ...RequestOption) Request {

client/rpc_client_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/micro/go-micro/errors"
 	"github.com/micro/go-micro/registry"
 	"github.com/micro/go-micro/registry/mock"
 	"github.com/micro/go-micro/selector"
@@ -56,6 +57,46 @@ func TestCallAddress(t *testing.T) {
 	}
 
 }
+
+func TestCallRetry(t *testing.T) {
+	service := "test.service"
+	method := "Test.Method"
+	address := "10.1.10.1:8080"
+
+	var called int
+
+	wrap := func(cf CallFunc) CallFunc {
+		return func(ctx context.Context, addr string, req Request, rsp interface{}, opts CallOptions) error {
+			called++
+			if called == 1 {
+				return errors.InternalServerError("test.error", "retry request")
+			}
+
+			// don't do the call
+			return nil
+		}
+	}
+
+	r := mock.NewRegistry()
+	c := NewClient(
+		Registry(r),
+		WrapCall(wrap),
+	)
+	c.Options().Selector.Init(selector.Registry(r))
+
+	req := c.NewRequest(service, method, nil)
+
+	// test calling remote address
+	if err := c.Call(context.Background(), req, nil, WithAddress(address)); err != nil {
+		t.Fatal("call with address error", err)
+	}
+
+	// num calls
+	if called < c.Options().CallOptions.Retries+1 {
+		t.Fatal("request not retried")
+	}
+}
+
 func TestCallWrapper(t *testing.T) {
 	var called bool
 	id := "test.1"

client/rpc_message.go

@@ -6,7 +6,16 @@ type message struct {
 	payload     interface{}
 }
 
-func newMessage(topic string, payload interface{}, contentType string) Message {
+func newMessage(topic string, payload interface{}, contentType string, opts ...MessageOption) Message {
+	var options MessageOptions
+	for _, o := range opts {
+		o(&options)
+	}
+
+	if len(options.ContentType) > 0 {
+		contentType = options.ContentType
+	}
+
 	return &message{
 		payload:     payload,
 		topic:       topic,

cmd/cmd.go

@@ -71,7 +71,6 @@ var (
 			Name:   "client_pool_size",
 			EnvVar: "MICRO_CLIENT_POOL_SIZE",
 			Usage:  "Sets the client connection pool size. Default: 1",
-			Value:  1,
 		},
 		cli.StringFlag{
 			Name:   "client_pool_ttl",
@@ -88,6 +87,11 @@ var (
 			EnvVar: "MICRO_REGISTER_INTERVAL",
 			Usage:  "Register interval in seconds",
 		},
+		cli.StringFlag{
+			Name:   "server",
+			EnvVar: "MICRO_SERVER",
+			Usage:  "Server for go-micro; rpc",
+		},
 		cli.StringFlag{
 			Name:   "server_name",
 			EnvVar: "MICRO_SERVER_NAME",
@@ -145,11 +149,6 @@ var (
 			Usage:  "Selector used to pick nodes for querying",
 			Value:  "cache",
 		},
-		cli.StringFlag{
-			Name:   "server",
-			EnvVar: "MICRO_SERVER",
-			Usage:  "Server for go-micro; rpc",
-		},
 		cli.StringFlag{
 			Name:   "transport",
 			EnvVar: "MICRO_TRANSPORT",
@@ -263,48 +262,40 @@ func (c *cmd) Before(ctx *cli.Context) error {
 
 	// Set the client
 	if name := ctx.String("client"); len(name) > 0 {
-		if cl, ok := c.opts.Clients[name]; ok {
+		// only change if we have the client and type differs
+		if cl, ok := c.opts.Clients[name]; ok && (*c.opts.Client).String() != name {
 			*c.opts.Client = cl()
 		}
 	}
 
 	// Set the server
 	if name := ctx.String("server"); len(name) > 0 {
-		if s, ok := c.opts.Servers[name]; ok {
+		// only change if we have the server and type differs
+		if s, ok := c.opts.Servers[name]; ok && (*c.opts.Server).String() != name {
 			*c.opts.Server = s()
 		}
 	}
 
 	// Set the broker
-	if name := ctx.String("broker"); len(name) > 0 || len(ctx.String("broker_address")) > 0 {
-		if len(name) == 0 {
-			name = defaultBroker
-		}
-
-		if b, ok := c.opts.Brokers[name]; ok {
-			n := b(broker.Addrs(strings.Split(ctx.String("broker_address"), ",")...))
-			*c.opts.Broker = n
-		} else {
+	if name := ctx.String("broker"); len(name) > 0 && (*c.opts.Broker).String() != name {
+		b, ok := c.opts.Brokers[name]
+		if !ok {
 			return fmt.Errorf("Broker %s not found", name)
 		}
 
+		*c.opts.Broker = b()
 		serverOpts = append(serverOpts, server.Broker(*c.opts.Broker))
 		clientOpts = append(clientOpts, client.Broker(*c.opts.Broker))
 	}
 
 	// Set the registry
-	if name := ctx.String("registry"); len(name) > 0 || len(ctx.String("registry_address")) > 0 {
-		if len(name) == 0 {
-			name = defaultRegistry
-		}
-
-		if r, ok := c.opts.Registries[name]; ok {
-			n := r(registry.Addrs(strings.Split(ctx.String("registry_address"), ",")...))
-			*c.opts.Registry = n
-		} else {
+	if name := ctx.String("registry"); len(name) > 0 && (*c.opts.Registry).String() != name {
+		r, ok := c.opts.Registries[name]
+		if !ok {
 			return fmt.Errorf("Registry %s not found", name)
 		}
 
+		*c.opts.Registry = r()
 		serverOpts = append(serverOpts, server.Registry(*c.opts.Registry))
 		clientOpts = append(clientOpts, client.Registry(*c.opts.Registry))
 
@@ -315,31 +306,26 @@ func (c *cmd) Before(ctx *cli.Context) error {
 	}
 
 	// Set the selector
-	if name := ctx.String("selector"); len(name) > 0 {
-		if s, ok := c.opts.Selectors[name]; ok {
-			n := s(selector.Registry(*c.opts.Registry))
-			*c.opts.Selector = n
-		} else {
+	if name := ctx.String("selector"); len(name) > 0 && (*c.opts.Selector).String() != name {
+		s, ok := c.opts.Selectors[name]
+		if !ok {
 			return fmt.Errorf("Selector %s not found", name)
 		}
 
+		*c.opts.Selector = s(selector.Registry(*c.opts.Registry))
+
 		// No server option here. Should there be?
 		clientOpts = append(clientOpts, client.Selector(*c.opts.Selector))
 	}
 
 	// Set the transport
-	if name := ctx.String("transport"); len(name) > 0 || len(ctx.String("transport_address")) > 0 {
-		if len(name) == 0 {
-			name = defaultTransport
-		}
-
-		if t, ok := c.opts.Transports[name]; ok {
-			n := t(transport.Addrs(strings.Split(ctx.String("transport_address"), ",")...))
-			*c.opts.Transport = n
-		} else {
+	if name := ctx.String("transport"); len(name) > 0 && (*c.opts.Transport).String() != name {
+		t, ok := c.opts.Transports[name]
+		if !ok {
 			return fmt.Errorf("Transport %s not found", name)
 		}
 
+		*c.opts.Transport = t()
 		serverOpts = append(serverOpts, server.Transport(*c.opts.Transport))
 		clientOpts = append(clientOpts, client.Transport(*c.opts.Transport))
 	}
@@ -360,6 +346,18 @@ func (c *cmd) Before(ctx *cli.Context) error {
 		serverOpts = append(serverOpts, server.Metadata(metadata))
 	}
 
+	if len(ctx.String("broker_address")) > 0 {
+		(*c.opts.Broker).Init(broker.Addrs(strings.Split(ctx.String("broker_address"), ",")...))
+	}
+
+	if len(ctx.String("registry_address")) > 0 {
+		(*c.opts.Registry).Init(registry.Addrs(strings.Split(ctx.String("registry_address"), ",")...))
+	}
+
+	if len(ctx.String("transport_address")) > 0 {
+		(*c.opts.Transport).Init(transport.Addrs(strings.Split(ctx.String("transport_address"), ",")...))
+	}
+
 	if len(ctx.String("server_name")) > 0 {
 		serverOpts = append(serverOpts, server.Name(ctx.String("server_name")))
 	}
@@ -385,7 +383,7 @@ func (c *cmd) Before(ctx *cli.Context) error {
 	}
 
 	// client opts
-	if r := ctx.Int("client_retries"); r > 0 {
+	if r := ctx.Int("client_retries"); r >= 0 {
 		clientOpts = append(clientOpts, client.Retries(r))
 	}
 

errors/errors.go

@@ -82,6 +82,16 @@ func NotFound(id, format string, a ...interface{}) error {
 	}
 }
 
+// MethodNotAllowed generates a 405 error.
+func MethodNotAllowed(id, format string, a ...interface{}) error {
+	return &Error{
+		Id:     id,
+		Code:   405,
+		Detail: fmt.Sprintf(format, a...),
+		Status: http.StatusText(405),
+	}
+}
+
 // InternalServerError generates a 500 error.
 func InternalServerError(id, format string, a ...interface{}) error {
 	return &Error{

function_test.go

@@ -15,8 +15,8 @@ func TestFunction(t *testing.T) {
 
 	// create service
 	fn := NewFunction(
-		Name("test.function"),
 		Registry(mock.NewRegistry()),
+		Name("test.function"),
 		AfterStart(func() error {
 			wg.Done()
 			return nil
@@ -26,29 +26,35 @@ func TestFunction(t *testing.T) {
 	// we can't test fn.Init as it parses the command line
 	// fn.Init()
 
+	ch := make(chan error, 2)
+
 	go func() {
-		// wait for start
-		wg.Wait()
-
-		// test call debug
-		req := fn.Client().NewRequest(
-			"test.function",
-			"Debug.Health",
-			new(proto.HealthRequest),
-		)
-
-		rsp := new(proto.HealthResponse)
-
-		err := fn.Client().Call(context.TODO(), req, rsp)
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		if rsp.Status != "ok" {
-			t.Fatalf("function response: %s", rsp.Status)
-		}
+		// run service
+		ch <- fn.Run()
 	}()
 
-	// run service
-	fn.Run()
+	// wait for start
+	wg.Wait()
+
+	// test call debug
+	req := fn.Client().NewRequest(
+		"test.function",
+		"Debug.Health",
+		new(proto.HealthRequest),
+	)
+
+	rsp := new(proto.HealthResponse)
+
+	err := fn.Client().Call(context.TODO(), req, rsp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if rsp.Status != "ok" {
+		t.Fatalf("function response: %s", rsp.Status)
+	}
+
+	if err := <-ch; err != nil {
+		t.Fatal(err)
+	}
 }

registry/consul/options.go

@@ -8,6 +8,16 @@ import (
 	"github.com/micro/go-micro/registry"
 )
 
+// Connect specifies services should be registered as Consul Connect services
+func Connect() registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, "consul_connect", true)
+	}
+}
+
 func Config(c *consul.Config) registry.Option {
 	return func(o *registry.Options) {
 		if o.Context == nil {

registry/consul_registry.go

@@ -19,10 +19,26 @@ type consulRegistry struct {
 	Client  *consul.Client
 	opts    Options
 
+	// connect enabled
+	connect bool
+
 	sync.Mutex
 	register map[string]uint64
 }
 
+func getDeregisterTTL(t time.Duration) time.Duration {
+	// splay slightly for the watcher?
+	splay := time.Second * 5
+	deregTTL := t + splay
+
+	// consul has a minimum timeout on deregistration of 1 minute.
+	if t < time.Minute {
+		deregTTL = time.Minute + splay
+	}
+
+	return deregTTL
+}
+
 func newTransport(config *tls.Config) *http.Transport {
 	if config == nil {
 		config = &tls.Config{
@@ -45,27 +61,31 @@ func newTransport(config *tls.Config) *http.Transport {
 	return t
 }
 
-func newConsulRegistry(opts ...Option) Registry {
-	var options Options
+func configure(c *consulRegistry, opts ...Option) {
+	// set opts
 	for _, o := range opts {
-		o(&options)
+		o(&c.opts)
 	}
 
 	// use default config
 	config := consul.DefaultConfig()
-	if options.Context != nil {
+
+	if c.opts.Context != nil {
 		// Use the consul config passed in the options, if available
-		if c, ok := options.Context.Value("consul_config").(*consul.Config); ok {
-			config = c
+		if co, ok := c.opts.Context.Value("consul_config").(*consul.Config); ok {
+			config = co
+		}
+		if cn, ok := c.opts.Context.Value("consul_connect").(bool); ok {
+			c.connect = cn
 		}
 	}
 
 	// check if there are any addrs
-	if len(options.Addrs) > 0 {
-		addr, port, err := net.SplitHostPort(options.Addrs[0])
+	if len(c.opts.Addrs) > 0 {
+		addr, port, err := net.SplitHostPort(c.opts.Addrs[0])
 		if ae, ok := err.(*net.AddrError); ok && ae.Err == "missing port in address" {
 			port = "8500"
-			addr = options.Addrs[0]
+			addr = c.opts.Addrs[0]
 			config.Address = fmt.Sprintf("%s:%s", addr, port)
 		} else if err == nil {
 			config.Address = fmt.Sprintf("%s:%s", addr, port)
@@ -73,34 +93,43 @@ func newConsulRegistry(opts ...Option) Registry {
 	}
 
 	// requires secure connection?
-	if options.Secure || options.TLSConfig != nil {
+	if c.opts.Secure || c.opts.TLSConfig != nil {
 		if config.HttpClient == nil {
 			config.HttpClient = new(http.Client)
 		}
 
 		config.Scheme = "https"
 		// We're going to support InsecureSkipVerify
-		config.HttpClient.Transport = newTransport(options.TLSConfig)
+		config.HttpClient.Transport = newTransport(c.opts.TLSConfig)
+	}
+
+	// set timeout
+	if c.opts.Timeout > 0 {
+		config.HttpClient.Timeout = c.opts.Timeout
 	}
 
 	// create the client
 	client, _ := consul.NewClient(config)
 
-	// set timeout
-	if options.Timeout > 0 {
-		config.HttpClient.Timeout = options.Timeout
-	}
+	// set address/client
+	c.Address = config.Address
+	c.Client = client
+}
 
+func newConsulRegistry(opts ...Option) Registry {
 	cr := &consulRegistry{
-		Address:  config.Address,
-		Client:   client,
-		opts:     options,
+		opts:     Options{},
 		register: make(map[string]uint64),
 	}
-
+	configure(cr, opts...)
 	return cr
 }
 
+func (c *consulRegistry) Init(opts ...Option) error {
+	configure(c, opts...)
+	return nil
+}
+
 func (c *consulRegistry) Deregister(s *Service) error {
 	if len(s.Nodes) == 0 {
 		return errors.New("Require at least one node")
@@ -115,19 +144,6 @@ func (c *consulRegistry) Deregister(s *Service) error {
 	return c.Client.Agent().ServiceDeregister(node.Id)
 }
 
-func getDeregisterTTL(t time.Duration) time.Duration {
-	// splay slightly for the watcher?
-	splay := time.Second * 5
-	deregTTL := t + splay
-
-	// consul has a minimum timeout on deregistration of 1 minute.
-	if t < time.Minute {
-		deregTTL = time.Minute + splay
-	}
-
-	return deregTTL
-}
-
 func (c *consulRegistry) Register(s *Service, opts ...RegisterOption) error {
 	if len(s.Nodes) == 0 {
 		return errors.New("Require at least one node")
@@ -198,14 +214,23 @@ func (c *consulRegistry) Register(s *Service, opts ...RegisterOption) error {
 	}
 
 	// register the service
-	if err := c.Client.Agent().ServiceRegister(&consul.AgentServiceRegistration{
+	asr := &consul.AgentServiceRegistration{
 		ID:      node.Id,
 		Name:    s.Name,
 		Tags:    tags,
 		Port:    node.Port,
 		Address: node.Address,
 		Check:   check,
-	}); err != nil {
+	}
+
+	// Specify consul connect
+	if c.connect {
+		asr.Connect = &consul.AgentServiceConnect{
+			Native: true,
+		}
+	}
+
+	if err := c.Client.Agent().ServiceRegister(asr); err != nil {
 		return err
 	}
 
@@ -224,7 +249,15 @@ func (c *consulRegistry) Register(s *Service, opts ...RegisterOption) error {
 }
 
 func (c *consulRegistry) GetService(name string) ([]*Service, error) {
-	rsp, _, err := c.Client.Health().Service(name, "", false, nil)
+	var rsp []*consul.ServiceEntry
+	var err error
+
+	// if we're connect enabled only get connect services
+	if c.connect {
+		rsp, _, err = c.Client.Health().Connect(name, "", false, nil)
+	} else {
+		rsp, _, err = c.Client.Health().Service(name, "", false, nil)
+	}
 	if err != nil {
 		return nil, err
 	}

registry/mdns/mdns.go

@@ -49,6 +49,17 @@ func newRegistry(opts ...registry.Option) registry.Registry {
 	}
 }
 
+func (m *mdnsRegistry) Init(opts ...registry.Option) error {
+	for _, o := range opts {
+		o(&m.opts)
+	}
+	return nil
+}
+
+func (m *mdnsRegistry) Options() registry.Options {
+	return m.opts
+}
+
 func (m *mdnsRegistry) Register(service *registry.Service, opts ...registry.RegisterOption) error {
 	m.Lock()
 	defer m.Unlock()
@@ -322,10 +333,6 @@ func (m *mdnsRegistry) String() string {
 	return "mdns"
 }
 
-func (m *mdnsRegistry) Options() registry.Options {
-	return m.opts
-}
-
 func NewRegistry(opts ...registry.Option) registry.Registry {
 	return newRegistry(opts...)
 }

registry/mock/mock.go

@@ -99,11 +99,15 @@ func (m *mockRegistry) String() string {
 	return "mock"
 }
 
+func (m *mockRegistry) Init(opts ...registry.Option) error {
+	return nil
+}
+
 func (m *mockRegistry) Options() registry.Options {
 	return registry.Options{}
 }
 
-func NewRegistry() registry.Registry {
+func NewRegistry(opts ...registry.Options) registry.Registry {
 	m := &mockRegistry{Services: make(map[string][]*registry.Service)}
 	m.init()
 	return m

registry/registry.go

@@ -9,13 +9,14 @@ import (
 // and an abstraction over varying implementations
 // {consul, etcd, zookeeper, ...}
 type Registry interface {
+	Init(...Option) error
+	Options() Options
 	Register(*Service, ...RegisterOption) error
 	Deregister(*Service) error
 	GetService(string) ([]*Service, error)
 	ListServices() ([]*Service, error)
 	Watch(...WatchOption) (Watcher, error)
 	String() string
-	Options() Options
 }
 
 type Option func(*Options)

server/rpc_server.go

@@ -66,6 +66,9 @@ func (s *rpcServer) accept(sock transport.Socket) {
 			return
 		}
 
+		// add to wait group
+		s.wg.Add(1)
+
 		// we use this Timeout header to set a server deadline
 		to := msg.Header["Timeout"]
 		// we use this Content-Type header to identify the codec needed
@@ -80,6 +83,7 @@ func (s *rpcServer) accept(sock transport.Socket) {
 				},
 				Body: []byte(err.Error()),
 			})
+			s.wg.Done()
 			return
 		}
 
@@ -102,15 +106,13 @@ func (s *rpcServer) accept(sock transport.Socket) {
 			}
 		}
 
-		// add to wait group
-		s.wg.Add(1)
-		defer s.wg.Done()
-
 		// TODO: needs better error handling
 		if err := s.rpc.serveRequest(ctx, codec, ct); err != nil {
+			s.wg.Done()
 			log.Logf("Unexpected error serving request, closing socket: %v", err)
 			return
 		}
+		s.wg.Done()
 	}
 }
 
@@ -386,6 +388,8 @@ func (s *rpcServer) Start() error {
 
 	log.Logf("Listening on %s", ts.Addr())
 	s.Lock()
+	// swap address
+	addr := s.opts.Address
 	s.opts.Address = ts.Addr()
 	s.Unlock()
 
@@ -405,6 +409,11 @@ func (s *rpcServer) Start() error {
 
 		// disconnect the broker
 		config.Broker.Disconnect()
+
+		s.Lock()
+		// swap back address
+		s.opts.Address = addr
+		s.Unlock()
 	}()
 
 	// TODO: subscribe to cruft

server/subscriber.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"strings"
 
 	"github.com/micro/go-micro/broker"
 	"github.com/micro/go-micro/codec"
@@ -176,6 +177,8 @@ func (s *rpcServer) createSubHandler(sb *subscriber, opts Options) broker.Handle
 		delete(hdr, "Content-Type")
 		ctx := metadata.NewContext(context.Background(), hdr)
 
+		results := make(chan error, len(sb.handlers))
+
 		for i := 0; i < len(sb.handlers); i++ {
 			handler := sb.handlers[i]
 
@@ -229,13 +232,26 @@ func (s *rpcServer) createSubHandler(sb *subscriber, opts Options) broker.Handle
 			s.wg.Add(1)
 			go func() {
 				defer s.wg.Done()
-				fn(ctx, &rpcMessage{
+				results <- fn(ctx, &rpcMessage{
 					topic:       sb.topic,
 					contentType: ct,
 					payload:     req.Interface(),
 				})
 			}()
 		}
+
+		var errors []string
+
+		for i := 0; i < len(sb.handlers); i++ {
+			if err := <-results; err != nil {
+				errors = append(errors, err.Error())
+			}
+		}
+
+		if len(errors) > 0 {
+			return fmt.Errorf("subscriber error: %s", strings.Join(errors, "\n"))
+		}
+
 		return nil
 	}
 }

service_test.go

@@ -31,29 +31,33 @@ func TestService(t *testing.T) {
 	// service.Init()
 
 	// run service
-	go service.Run()
+	go func() {
+		// wait for start
+		wg.Wait()
 
-	// wait for start
-	wg.Wait()
+		// test call debug
+		req := service.Client().NewRequest(
+			"test.service",
+			"Debug.Health",
+			new(proto.HealthRequest),
+		)
 
-	// test call debug
-	req := service.Client().NewRequest(
-		"test.service",
-		"Debug.Health",
-		new(proto.HealthRequest),
-	)
+		rsp := new(proto.HealthResponse)
 
-	rsp := new(proto.HealthResponse)
+		err := service.Client().Call(context.TODO(), req, rsp)
+		if err != nil {
+			t.Fatal(err)
+		}
 
-	err := service.Client().Call(context.TODO(), req, rsp)
-	if err != nil {
+		if rsp.Status != "ok" {
+			t.Fatalf("service response: %s", rsp.Status)
+		}
+
+		// shutdown the service
+		cancel()
+	}()
+
+	if err := service.Run(); err != nil {
 		t.Fatal(err)
 	}
-
-	if rsp.Status != "ok" {
-		t.Fatalf("service response: %s", rsp.Status)
-	}
-
-	// shutdown the service
-	cancel()
 }

transport/http_transport.go

@@ -423,6 +423,17 @@ func (h *httpTransport) Listen(addr string, opts ...ListenOption) (Listener, err
 	}, nil
 }
 
+func (h *httpTransport) Init(opts ...Option) error {
+	for _, o := range opts {
+		o(&h.opts)
+	}
+	return nil
+}
+
+func (h *httpTransport) Options() Options {
+	return h.opts
+}
+
 func (h *httpTransport) String() string {
 	return "http"
 }

transport/mock/mock.go

@@ -171,6 +171,17 @@ func (m *mockTransport) Listen(addr string, opts ...transport.ListenOption) (tra
 	return listener, nil
 }
 
+func (m *mockTransport) Init(opts ...transport.Option) error {
+	for _, o := range opts {
+		o(&m.opts)
+	}
+	return nil
+}
+
+func (m *mockTransport) Options() transport.Options {
+	return m.opts
+}
+
 func (m *mockTransport) String() string {
 	return "mock"
 }

transport/transport.go

@@ -30,6 +30,8 @@ type Listener interface {
 // services. It uses socket send/recv semantics and had various
 // implementations {HTTP, RabbitMQ, NATS, ...}
 type Transport interface {
+	Init(...Option) error
+	Options() Options
 	Dial(addr string, opts ...DialOption) (Client, error)
 	Listen(addr string, opts ...ListenOption) (Listener, error)
 	String() string