Use registry.options.Context to set Consul TCP check option

registry/consul/options.go

@@ -2,6 +2,7 @@ package consul
 
 import (
 	"context"
+	"time"
 
 	consul "github.com/hashicorp/consul/api"
 	"github.com/micro/go-micro/registry"
@@ -15,3 +16,23 @@ func Config(c *consul.Config) registry.Option {
 		o.Context = context.WithValue(o.Context, "consul_config", c)
 	}
 }
+
+//
+// RegisterTCPCheck will tell the service provider to check the service address
+// and port every `t` interval. It will enabled only if `t` is greater than 0.
+// See `TCP + Interval` for more information [1].
+//
+// [1] https://www.consul.io/docs/agent/checks.html
+//
+func RegisterTCPCheck(t time.Duration) registry.Option {
+	return func(o *registry.Options) {
+		if t <= time.Duration(0) {
+			return
+		}
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context,
+			registry.ConsulRegisterTCPCheckKey, t)
+	}
+}

registry/consul_registry.go

@@ -14,6 +14,10 @@ import (
 	hash "github.com/mitchellh/hashstructure"
 )
 
+const (
+	ConsulRegisterTCPCheckKey = "consul_register_tcp_check"
+)
+
 type consulRegistry struct {
 	Address string
 	Client  *consul.Client
@@ -138,6 +142,15 @@ func (c *consulRegistry) Register(s *Service, opts ...RegisterOption) error {
 		o(&options)
 	}
 
+	if c.opts.Context != nil {
+		tcpCheckInterval, ok := c.opts.Context.
+			Value(ConsulRegisterTCPCheckKey).(time.Duration)
+		if ok {
+			options.TCPCheck = true
+			options.Interval = tcpCheckInterval
+		}
+	}
+
 	// create hash of service; uint64
 	h, err := hash.Hash(s, nil)
 	if err != nil {

registry/options.go

@@ -68,23 +68,6 @@ func RegisterTTL(t time.Duration) RegisterOption {
 	}
 }
 
-//
-// RegisterTCPCheck will tell the service provider to check the service address
-// and port every `t` interval. It will enabled only if `t` is greater than 0.
-// This option is for registry using Consul, see `TCP + Interval` more
-// information [1].
-//
-// [1] https://www.consul.io/docs/agent/checks.html
-//
-func RegisterTCPCheck(t time.Duration) RegisterOption {
-	return func(o *RegisterOptions) {
-		if t > time.Duration(0) {
-			o.TCPCheck = true
-			o.Interval = t
-		}
-	}
-}
-
 // Watch a service
 func WatchService(name string) WatchOption {
 	return func(o *WatchOptions) {

server/options.go

@@ -25,9 +25,7 @@ type Options struct {
 	HdlrWrappers []HandlerWrapper
 	SubWrappers  []SubscriberWrapper
 
-	RegisterTCPCheck bool
-	RegisterTTL      time.Duration
-	RegisterInterval time.Duration
+	RegisterTTL time.Duration
 
 	// Debug Handler which can be set by a user
 	DebugHandler debug.DebugHandler
@@ -166,23 +164,6 @@ func RegisterTTL(t time.Duration) Option {
 	}
 }
 
-//
-// RegisterTCPCheck will tell the service provider to check the service address
-// and port every `t` interval. It will enabled only if `t` is greater than 0.
-// This option is for registry using Consul, see `TCP + Interval` more
-// information [1].
-//
-// [1] https://www.consul.io/docs/agent/checks.html
-//
-func RegisterTCPCheck(t time.Duration) Option {
-	return func(o *Options) {
-		if t > time.Duration(0) {
-			o.RegisterTCPCheck = true
-			o.RegisterInterval = t
-		}
-	}
-}
-
 // Wait tells the server to wait for requests to finish before exiting
 func Wait(b bool) Option {
 	return func(o *Options) {

server/rpc_server.go

@@ -278,10 +278,7 @@ func (s *rpcServer) Register() error {
 	}
 
 	// create registry options
-	rOpts := []registry.RegisterOption{
-		registry.RegisterTTL(config.RegisterTTL),
-		registry.RegisterTCPCheck(config.RegisterInterval),
-	}
+	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)}
 
 	if err := config.Registry.Register(service, rOpts...); err != nil {
 		return err