Skip to content

Commit

Permalink
reverse_proxy: Fix random_choose lb_policy
Browse files Browse the repository at this point in the history
  • Loading branch information
d-masson committed Oct 19, 2020
1 parent 385adf5 commit 24ae784
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 7 deletions.
16 changes: 9 additions & 7 deletions modules/caddyhttp/reverseproxy/selectionpolicies.go
Original file line number Diff line number Diff line change
Expand Up @@ -139,14 +139,17 @@ func (r RandomChoiceSelection) Select(pool UpstreamPool, _ *http.Request) *Upstr
if k > len(pool) {
k = len(pool)
}
choices := make([]*Upstream, k)
// Slice of length 0 and capacity k (if we use an array and have some unavailable
// upstream, we could have some empty element which may lead
// to an index out of range runtime error)
choices := make([]*Upstream, 0, k)
for i, upstream := range pool {
if !upstream.Available() {
continue
}
j := weakrand.Intn(i + 1)
if j < k {
choices[j] = upstream
choices = append(choices, upstream)
}
}
return leastRequests(choices)
Expand Down Expand Up @@ -397,13 +400,12 @@ func leastRequests(upstreams []*Upstream) *Upstream {
return nil
}
var best []*Upstream
var bestReqs int
var bestReqs int = -1
for _, upstream := range upstreams {
reqs := upstream.NumRequests()
if reqs == 0 {
return upstream
}
if reqs <= bestReqs {
// If bestReqs was just initialized to -1
// we need to append upstream also
if reqs <= bestReqs || bestReqs == -1 {
bestReqs = reqs
best = append(best, upstream)
}
Expand Down
51 changes: 51 additions & 0 deletions modules/caddyhttp/reverseproxy/selectionpolicies_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -271,3 +271,54 @@ func TestURIHashPolicy(t *testing.T) {
t.Error("Expected uri policy policy host to be nil.")
}
}

func TestLeastRequests(t *testing.T) {
pool := testPool()
pool[0].Dial = "localhost:8080"
pool[1].Dial = "localhost:8081"
pool[2].Dial = "localhost:8082"
pool[0].SetHealthy(true)
pool[1].SetHealthy(true)
pool[2].SetHealthy(true)
pool[0].CountRequest(10)
pool[1].CountRequest(20)
pool[2].CountRequest(30)

result := leastRequests(pool)

if result == nil {
t.Error("Least request should not return nil")
}

if result != pool[0] {
t.Error("Least request should return pool[0]")
}
}

func TestRandomChoicePolicy(t *testing.T) {
pool := testPool()
pool[0].Dial = "localhost:8080"
pool[1].Dial = "localhost:8081"
pool[2].Dial = "localhost:8082"
pool[0].SetHealthy(false)
pool[1].SetHealthy(true)
pool[2].SetHealthy(true)
pool[0].CountRequest(10)
pool[1].CountRequest(20)
pool[2].CountRequest(30)

request := httptest.NewRequest(http.MethodGet, "/test", nil)
randomChoicePolicy := new(RandomChoiceSelection)
randomChoicePolicy.Choose = 2

h := randomChoicePolicy.Select(pool, request)

if h == nil {
t.Error("RandomChoicePolicy should not return nil")
}

if h != pool[1] {
t.Error("RandomChoicePolicy should choose pool[1]")
}

}

0 comments on commit 24ae784

Please sign in to comment.