From 40ab1f325ccf4630449cb6f2e5ad5009bcd8218c Mon Sep 17 00:00:00 2001
From: Douglas Wagner <DougWagner@users.noreply.github.com>
Date: Wed, 7 Aug 2019 08:14:04 -0700
Subject: [PATCH] Wrr loadbalancer honors old weight on recovered servers

---
 healthcheck/healthcheck.go      | 36 +++++++++++++++++++++++----------
 healthcheck/healthcheck_test.go |  2 +-
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/healthcheck/healthcheck.go b/healthcheck/healthcheck.go
index 7d5b1d864..7c7d2851f 100644
--- a/healthcheck/healthcheck.go
+++ b/healthcheck/healthcheck.go
@@ -49,11 +49,16 @@ func (opt Options) String() string {
 	return fmt.Sprintf("[Hostname: %s Headers: %v Path: %s Port: %d Interval: %s]", opt.Hostname, opt.Headers, opt.Path, opt.Port, opt.Interval)
 }
 
+type backendURL struct {
+	url    *url.URL
+	weight int
+}
+
 // BackendConfig HealthCheck configuration for a backend
 type BackendConfig struct {
 	Options
 	name           string
-	disabledURLs   []*url.URL
+	disabledURLs   []backendURL
 	requestTimeout time.Duration
 }
 
@@ -129,18 +134,18 @@ func (hc *HealthCheck) execute(ctx context.Context, backend *BackendConfig) {
 
 func (hc *HealthCheck) checkBackend(backend *BackendConfig) {
 	enabledURLs := backend.LB.Servers()
-	var newDisabledURLs []*url.URL
-	for _, url := range backend.disabledURLs {
+	var newDisabledURLs []backendURL
+	for _, backendurl := range backend.disabledURLs {
 		serverUpMetricValue := float64(0)
-		if err := checkHealth(url, backend); err == nil {
-			log.Warnf("Health check up: Returning to server list. Backend: %q URL: %q", backend.name, url.String())
-			backend.LB.UpsertServer(url, roundrobin.Weight(1))
+		if err := checkHealth(backendurl.url, backend); err == nil {
+			log.Warnf("Health check up: Returning to server list. Backend: %q URL: %q Weight: %d", backend.name, backendurl.url.String(), backendurl.weight)
+			backend.LB.UpsertServer(backendurl.url, roundrobin.Weight(backendurl.weight))
 			serverUpMetricValue = 1
 		} else {
-			log.Warnf("Health check still failing. Backend: %q URL: %q Reason: %s", backend.name, url.String(), err)
-			newDisabledURLs = append(newDisabledURLs, url)
+			log.Warnf("Health check still failing. Backend: %q URL: %q Reason: %s", backend.name, backendurl.url.String(), err)
+			newDisabledURLs = append(newDisabledURLs, backendurl)
 		}
-		labelValues := []string{"backend", backend.name, "url", url.String()}
+		labelValues := []string{"backend", backend.name, "url", backendurl.url.String()}
 		hc.metrics.BackendServerUpGauge().With(labelValues...).Set(serverUpMetricValue)
 	}
 	backend.disabledURLs = newDisabledURLs
@@ -148,9 +153,18 @@ func (hc *HealthCheck) checkBackend(backend *BackendConfig) {
 	for _, url := range enabledURLs {
 		serverUpMetricValue := float64(1)
 		if err := checkHealth(url, backend); err != nil {
-			log.Warnf("Health check failed: Remove from server list. Backend: %q URL: %q Reason: %s", backend.name, url.String(), err)
+			weight := 1
+			rr, ok := backend.LB.(*roundrobin.RoundRobin)
+			if ok {
+				var gotWeight bool
+				weight, gotWeight = rr.ServerWeight(url)
+				if !gotWeight {
+					weight = 1
+				}
+			}
+			log.Warnf("Health check failed: Remove from server list. Backend: %q URL: %q Weight: %d Reason: %s", backend.name, url.String(), weight, err)
 			backend.LB.RemoveServer(url)
-			backend.disabledURLs = append(backend.disabledURLs, url)
+			backend.disabledURLs = append(backend.disabledURLs, backendURL{url, weight})
 			serverUpMetricValue = 0
 		}
 		labelValues := []string{"backend", backend.name, "url", url.String()}
diff --git a/healthcheck/healthcheck_test.go b/healthcheck/healthcheck_test.go
index 1b6b4c3c1..0cc466d7f 100644
--- a/healthcheck/healthcheck_test.go
+++ b/healthcheck/healthcheck_test.go
@@ -112,7 +112,7 @@ func TestSetBackendsConfiguration(t *testing.T) {
 			if test.startHealthy {
 				lb.servers = append(lb.servers, serverURL)
 			} else {
-				backend.disabledURLs = append(backend.disabledURLs, serverURL)
+				backend.disabledURLs = append(backend.disabledURLs, backendURL{serverURL, 1})
 			}
 
 			collectingMetrics := testhelpers.NewCollectingHealthCheckMetrics()