From 361946eb0c08791ad16ebc3e82a79512895e650f Mon Sep 17 00:00:00 2001
From: Saber Haj Rabiee <sabify@users.noreply.github.com>
Date: Tue, 20 Jun 2023 10:42:58 -0700
Subject: reverseproxy: weighted_round_robin load balancing policy (#5579)

* added weighted round robin algorithm to load balancer

* added an adapt integration test for wrr and fixed a typo

* changed args format to Caddyfile args convention

* added provisioner and validator for wrr

* simplified the code and improved doc
---
 .../reverseproxy/selectionpolicies_test.go         | 57 ++++++++++++++++++++++
 1 file changed, 57 insertions(+)

(limited to 'modules/caddyhttp/reverseproxy/selectionpolicies_test.go')

diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
index 93dcb77..dc613a5 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
@@ -74,6 +74,63 @@ func TestRoundRobinPolicy(t *testing.T) {
 	}
 }
 
+func TestWeightedRoundRobinPolicy(t *testing.T) {
+	pool := testPool()
+	wrrPolicy := WeightedRoundRobinSelection{
+		Weights:     []int{3, 2, 1},
+		totalWeight: 6,
+	}
+	req, _ := http.NewRequest("GET", "/", nil)
+
+	h := wrrPolicy.Select(pool, req, nil)
+	if h != pool[0] {
+		t.Error("Expected first weighted round robin host to be first host in the pool.")
+	}
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[0] {
+		t.Error("Expected second weighted round robin host to be first host in the pool.")
+	}
+	// Third selected host is 1, because counter starts at 0
+	// and increments before host is selected
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[1] {
+		t.Error("Expected third weighted round robin host to be second host in the pool.")
+	}
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[1] {
+		t.Error("Expected fourth weighted round robin host to be second host in the pool.")
+	}
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[2] {
+		t.Error("Expected fifth weighted round robin host to be third host in the pool.")
+	}
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[0] {
+		t.Error("Expected sixth weighted round robin host to be first host in the pool.")
+	}
+
+	// mark host as down
+	pool[0].setHealthy(false)
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[1] {
+		t.Error("Expected to skip down host.")
+	}
+	// mark host as up
+	pool[0].setHealthy(true)
+
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[0] {
+		t.Error("Expected to select first host on availablity.")
+	}
+	// mark host as full
+	pool[1].countRequest(1)
+	pool[1].MaxRequests = 1
+	h = wrrPolicy.Select(pool, req, nil)
+	if h != pool[2] {
+		t.Error("Expected to skip full host.")
+	}
+}
+
 func TestLeastConnPolicy(t *testing.T) {
 	pool := testPool()
 	lcPolicy := LeastConnSelection{}
-- 
cgit v1.2.3