Use registry.options.Context to set Consul TCP check option
This commit is contained in:
		| @@ -2,6 +2,7 @@ package consul | |||||||
|  |  | ||||||
| import ( | import ( | ||||||
| 	"context" | 	"context" | ||||||
|  | 	"time" | ||||||
|  |  | ||||||
| 	consul "github.com/hashicorp/consul/api" | 	consul "github.com/hashicorp/consul/api" | ||||||
| 	"github.com/micro/go-micro/registry" | 	"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) | 		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) | ||||||
|  | 	} | ||||||
|  | } | ||||||
|   | |||||||
| @@ -14,6 +14,10 @@ import ( | |||||||
| 	hash "github.com/mitchellh/hashstructure" | 	hash "github.com/mitchellh/hashstructure" | ||||||
| ) | ) | ||||||
|  |  | ||||||
|  | const ( | ||||||
|  | 	ConsulRegisterTCPCheckKey = "consul_register_tcp_check" | ||||||
|  | ) | ||||||
|  |  | ||||||
| type consulRegistry struct { | type consulRegistry struct { | ||||||
| 	Address string | 	Address string | ||||||
| 	Client  *consul.Client | 	Client  *consul.Client | ||||||
| @@ -138,6 +142,15 @@ func (c *consulRegistry) Register(s *Service, opts ...RegisterOption) error { | |||||||
| 		o(&options) | 		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 | 	// create hash of service; uint64 | ||||||
| 	h, err := hash.Hash(s, nil) | 	h, err := hash.Hash(s, nil) | ||||||
| 	if err != nil { | 	if err != nil { | ||||||
|   | |||||||
| @@ -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 | // Watch a service | ||||||
| func WatchService(name string) WatchOption { | func WatchService(name string) WatchOption { | ||||||
| 	return func(o *WatchOptions) { | 	return func(o *WatchOptions) { | ||||||
|   | |||||||
| @@ -25,9 +25,7 @@ type Options struct { | |||||||
| 	HdlrWrappers []HandlerWrapper | 	HdlrWrappers []HandlerWrapper | ||||||
| 	SubWrappers  []SubscriberWrapper | 	SubWrappers  []SubscriberWrapper | ||||||
|  |  | ||||||
| 	RegisterTCPCheck bool | 	RegisterTTL time.Duration | ||||||
| 	RegisterTTL      time.Duration |  | ||||||
| 	RegisterInterval time.Duration |  | ||||||
|  |  | ||||||
| 	// Debug Handler which can be set by a user | 	// Debug Handler which can be set by a user | ||||||
| 	DebugHandler debug.DebugHandler | 	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 | // Wait tells the server to wait for requests to finish before exiting | ||||||
| func Wait(b bool) Option { | func Wait(b bool) Option { | ||||||
| 	return func(o *Options) { | 	return func(o *Options) { | ||||||
|   | |||||||
| @@ -278,10 +278,7 @@ func (s *rpcServer) Register() error { | |||||||
| 	} | 	} | ||||||
|  |  | ||||||
| 	// create registry options | 	// create registry options | ||||||
| 	rOpts := []registry.RegisterOption{ | 	rOpts := []registry.RegisterOption{registry.RegisterTTL(config.RegisterTTL)} | ||||||
| 		registry.RegisterTTL(config.RegisterTTL), |  | ||||||
| 		registry.RegisterTCPCheck(config.RegisterInterval), |  | ||||||
| 	} |  | ||||||
|  |  | ||||||
| 	if err := config.Registry.Register(service, rOpts...); err != nil { | 	if err := config.Registry.Register(service, rOpts...); err != nil { | ||||||
| 		return err | 		return err | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user