From fee4890e947f2b3c115c416015a7097da1817692 Mon Sep 17 00:00:00 2001
From: Andrew Hamon <andrew@hamon.cc>
Date: Tue, 14 Jun 2016 17:43:06 -0400
Subject: [PATCH] Balance round robin evenly when some hosts are down (#880)

* Balance round robin evenly when some hosts are down

Before, when load balancing across multiple hosts, if a host went down
then the next host in line would be sent a double share of requests.
This is because the round robin counter was only incremented once per
request, regardless of the health of the selection. If current
selection was unhealthy then the policy would advance to the next host,
but this would not be reflected in the policy counter. To fix this, the
counter is now incremented for every attempted host.

This commit adds a test case that identifies the issue, and a fix.

* Make robin counter private

* Use a mutex to sync round robin selection
---
 caddyhttp/proxy/policy.go      | 24 +++++++++++++-----------
 caddyhttp/proxy/policy_test.go | 15 +++++++++++----
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/caddyhttp/proxy/policy.go b/caddyhttp/proxy/policy.go
index 96e382a5c..0bf657cae 100644
--- a/caddyhttp/proxy/policy.go
+++ b/caddyhttp/proxy/policy.go
@@ -2,7 +2,7 @@ package proxy
 
 import (
 	"math/rand"
-	"sync/atomic"
+	"sync"
 )
 
 // HostPool is a collection of UpstreamHosts.
@@ -82,20 +82,22 @@ func (r *LeastConn) Select(pool HostPool) *UpstreamHost {
 
 // RoundRobin is a policy that selects hosts based on round robin ordering.
 type RoundRobin struct {
-	Robin uint32
+	robin uint32
+	mutex sync.Mutex
 }
 
 // Select selects an up host from the pool using a round robin ordering scheme.
 func (r *RoundRobin) Select(pool HostPool) *UpstreamHost {
 	poolLen := uint32(len(pool))
-	selection := atomic.AddUint32(&r.Robin, 1) % poolLen
-	host := pool[selection]
-	// if the currently selected host is not available, just ffwd to up host
-	for i := uint32(1); !host.Available() && i < poolLen; i++ {
-		host = pool[(selection+i)%poolLen]
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+	// Return next available host
+	for i := uint32(0); i < poolLen; i++ {
+		r.robin++
+		host := pool[r.robin%poolLen]
+		if host.Available() {
+			return host
+		}
 	}
-	if !host.Available() {
-		return nil
-	}
-	return host
+	return nil
 }
diff --git a/caddyhttp/proxy/policy_test.go b/caddyhttp/proxy/policy_test.go
index 4cc05f029..4ef591195 100644
--- a/caddyhttp/proxy/policy_test.go
+++ b/caddyhttp/proxy/policy_test.go
@@ -63,11 +63,18 @@ func TestRoundRobinPolicy(t *testing.T) {
 	if h != pool[2] {
 		t.Error("Expected to skip down host.")
 	}
-	// mark host as full
-	pool[2].Conns = 1
-	pool[2].MaxConns = 1
+	// mark host as up
+	pool[1].Unhealthy = false
+
 	h = rrPolicy.Select(pool)
-	if h != pool[0] {
+	if h == pool[2] {
+		t.Error("Expected to balance evenly among healthy hosts")
+	}
+	// mark host as full
+	pool[1].Conns = 1
+	pool[1].MaxConns = 1
+	h = rrPolicy.Select(pool)
+	if h != pool[2] {
 		t.Error("Expected to skip full host.")
 	}
 }