From d7b3765c71ceb13fac631aa2ad7d37505a9db0dc Mon Sep 17 00:00:00 2001
From: Asim <asim@chuhnk.me>
Date: Sat, 19 Dec 2015 18:28:08 +0000
Subject: [PATCH] Allow setting of timeout for registry

---
 registry/consul_registry.go | 19 ++++++++++++++++++-
 registry/options.go         | 15 +++++++++++++++
 registry/registry.go        |  4 +---
 3 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 registry/options.go

diff --git a/registry/consul_registry.go b/registry/consul_registry.go
index 0a04d00f..ea287642 100644
--- a/registry/consul_registry.go
+++ b/registry/consul_registry.go
@@ -12,6 +12,7 @@ import (
 type consulRegistry struct {
 	Address string
 	Client  *consul.Client
+	Options Options
 }
 
 func encodeEndpoints(en []*Endpoint) []string {
@@ -69,7 +70,20 @@ func decodeMetadata(tags []string) map[string]string {
 }
 
 func newConsulRegistry(addrs []string, opts ...Option) Registry {
+	var opt Options
+	for _, o := range opts {
+		o(&opt)
+	}
+
+	// use default config
 	config := consul.DefaultConfig()
+
+	// set timeout
+	if opt.Timeout > 0 {
+		config.HttpClient.Timeout = opt.Timeout
+	}
+
+	// check if there are any addrs
 	if len(addrs) > 0 {
 		addr, port, err := net.SplitHostPort(addrs[0])
 		if ae, ok := err.(*net.AddrError); ok && ae.Err == "missing port in address" {
@@ -79,11 +93,14 @@ func newConsulRegistry(addrs []string, opts ...Option) Registry {
 			config.Address = fmt.Sprintf("%s:%s", addr, port)
 		}
 	}
+
+	// create the client
 	client, _ := consul.NewClient(config)
 
 	cr := &consulRegistry{
 		Address: config.Address,
 		Client:  client,
+		Options: opt,
 	}
 
 	return cr
@@ -178,7 +195,7 @@ func (c *consulRegistry) GetService(name string) ([]*Service, error) {
 }
 
 func (c *consulRegistry) ListServices() ([]*Service, error) {
-	rsp, _, err := c.Client.Catalog().Services(&consul.QueryOptions{})
+	rsp, _, err := c.Client.Catalog().Services(nil)
 	if err != nil {
 		return nil, err
 	}
diff --git a/registry/options.go b/registry/options.go
new file mode 100644
index 00000000..fbd16b34
--- /dev/null
+++ b/registry/options.go
@@ -0,0 +1,15 @@
+package registry
+
+import (
+	"time"
+)
+
+type Options struct {
+	Timeout time.Duration
+}
+
+func Timeout(t time.Duration) Option {
+	return func(o *Options) {
+		o.Timeout = t
+	}
+}
diff --git a/registry/registry.go b/registry/registry.go
index 4f211d6e..0f139993 100644
--- a/registry/registry.go
+++ b/registry/registry.go
@@ -8,9 +8,7 @@ type Registry interface {
 	Watch() (Watcher, error)
 }
 
-type options struct{}
-
-type Option func(*options)
+type Option func(*Options)
 
 var (
 	DefaultRegistry = newConsulRegistry([]string{})