From 980406db7dedd372c23d640590dd7535a0776e19 Mon Sep 17 00:00:00 2001
From: Afshin Darian <git@darian.af>
Date: Wed, 28 Dec 2016 16:13:06 +0000
Subject: [PATCH] Refactor client `has` function.

---
 client.go      | 32 +++++++++++++++++---------------
 error.go       |  1 +
 sleuth_test.go | 10 ++++++----
 3 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/client.go b/client.go
index 8f7bb37..74ef48e 100644
--- a/client.go
+++ b/client.go
@@ -63,17 +63,17 @@ func (c *Client) add(group, name, node, service, version string) error {
 
 // Blocks until the required services are available to the client.
 // Returns true if it had to block and false if it returns immediately.
-func (c *Client) block(services ...string) bool {
+func (c *Client) block(required map[string]struct{}, services []string) bool {
 	// Even though the client may have just checked to see if services exist,
 	// the check is performed here in case there was a delay waiting for the
 	// additions mutex to become available.
-	if c.has(services...) {
+	if c.has(required) {
 		return false
 	}
 	c.log.Blocked("sleuth: waiting for client to find %s", services)
 	c.additions.activate()
 	for range c.additions.stream {
-		if c.has(services...) {
+		if c.has(required) {
 			break
 		}
 	}
@@ -161,21 +161,15 @@ func (c *Client) Do(req *http.Request) (*http.Response, error) {
 	return nil, newError(errTimeout, "%s {%s}%s timed out", req.Method, to, url)
 }
 
-func (c *Client) has(services ...string) bool {
-	// Check to see if required services are already registered.
-	verified := make(map[string]bool)
+func (c *Client) has(required map[string]struct{}) bool {
+	// Check to see if required services already exist locally.
 	available := 0
-	for _, service := range services {
-		verified[service] = false
-	}
-	total := len(verified)
-	for service := range verified {
+	for service := range required {
 		if peers, ok := c.services.get(service); ok && peers.available() {
-			verified[service] = true
 			available += 1
 		}
 	}
-	return available == total
+	return available == len(required)
 }
 
 func (c *Client) listen(handle string, listener chan *http.Response) {
@@ -236,8 +230,16 @@ func (c *Client) WaitFor(services ...string) error {
 	if c.closed {
 		return newError(errClosed, "client is closed").escalate(errWait)
 	}
-	if !c.has(services...) {
-		c.block(services...)
+	// Collapse services and make sure all values are unique.
+	required := make(map[string]struct{})
+	for _, service := range services {
+		required[service] = struct{}{}
+	}
+	if len(required) != len(services) {
+		c.log.Warn("sleuth: %v contains duplicates [%d]", services, warnDuplicate)
+	}
+	if !c.has(required) {
+		c.block(required, services)
 	}
 	return nil
 }
diff --git a/error.go b/error.go
index 3039d9a..107c259 100644
--- a/error.go
+++ b/error.go
@@ -10,6 +10,7 @@ const (
 	// Warnings are in the 801-899 range.
 	warnInterface = 801
 	warnClose     = 802
+	warnDuplicate = 803
 	// Errors are in the 901-999 range.
 	errNew              = 901
 	errDispatch         = 902
diff --git a/sleuth_test.go b/sleuth_test.go
index 983fcc0..64b724d 100644
--- a/sleuth_test.go
+++ b/sleuth_test.go
@@ -426,7 +426,7 @@ func TestZipUnzip(t *testing.T) {
 
 func TestIntegratedCycle(t *testing.T) {
 	addr := "sleuth-test-server-one"
-	client, err := New(&Config{group: GROUP})
+	client, err := New(&Config{group: GROUP, LogLevel: "warn"})
 	if err != nil {
 		t.Errorf("client instantiation failed: %s", err.Error())
 		return
@@ -450,11 +450,13 @@ func TestIntegratedCycle(t *testing.T) {
 			t.Errorf("server close failed: %s", err.Error())
 		}
 	}(server, t)
-	// Wait until the server becomes available.
-	client.WaitFor(addr)
+	// Wait until the server becomes available. Confirm dupicate warning.
+	client.WaitFor(addr, addr)
 	// Set timeout to 10 seconds to accommodate slow test spin-up.
 	client.Timeout = time.Second * 10
-	if client.block(addr) {
+	required := make(map[string]struct{})
+	required[addr] = struct{}{}
+	if client.block(required, []string{addr}) {
 		t.Errorf("call to block should have returned immediately")
 	}
 	body := "foo bar baz"