From 14075e31dca154380a2ae54272f1221a2c739ff0 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Wed, 21 Feb 2018 00:14:18 +0100
Subject: [PATCH] Vastly improve the way load balancing works

---
 dnscrypt-proxy/common.go      | 14 ++++++++++++++
 dnscrypt-proxy/serversInfo.go | 31 ++++++++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/dnscrypt-proxy/common.go b/dnscrypt-proxy/common.go
index 8dadfc84..2793c5fb 100644
--- a/dnscrypt-proxy/common.go
+++ b/dnscrypt-proxy/common.go
@@ -83,6 +83,20 @@ func Max(a, b int) int {
 	return b
 }
 
+func MinF(a, b float64) float64 {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+func MaxF(a, b float64) float64 {
+	if a > b {
+		return a
+	}
+	return b
+}
+
 func StringReverse(s string) string {
 	r := []rune(s)
 	for i, j := 0, len(r)-1; i < len(r)/2; i, j = i+1, j-1 {
diff --git a/dnscrypt-proxy/serversInfo.go b/dnscrypt-proxy/serversInfo.go
index 7ffab692..a2574843 100644
--- a/dnscrypt-proxy/serversInfo.go
+++ b/dnscrypt-proxy/serversInfo.go
@@ -167,8 +167,28 @@ func (serversInfo *ServersInfo) getOne() *ServerInfo {
 	if candidate == 0 {
 		return &serversInfo.inner[candidate]
 	}
-	if serversInfo.inner[candidate].rtt.Value() < serversInfo.inner[0].rtt.Value() {
+	candidateRtt, currentBestRtt := serversInfo.inner[candidate].rtt.Value(), serversInfo.inner[0].rtt.Value()
+	if currentBestRtt < 0 {
+		currentBestRtt = candidateRtt
+		serversInfo.inner[0].rtt.Set(currentBestRtt)
+	}
+	partialSort := false
+	if candidateRtt < currentBestRtt {
 		serversInfo.inner[candidate], serversInfo.inner[0] = serversInfo.inner[0], serversInfo.inner[candidate]
+		partialSort = true
+		dlog.Debugf("New prefered candidate: %v (rtt: %v vs previous: %v)", serversInfo.inner[0].Name, candidateRtt, currentBestRtt)
+	} else if candidateRtt >= currentBestRtt*4.0 {
+		if time.Since(serversInfo.inner[candidate].lastActionTS) > time.Duration(1*time.Minute) {
+			serversInfo.inner[candidate].rtt.Add(MinF(MaxF(candidateRtt/2.0, currentBestRtt*2.0), candidateRtt))
+			partialSort = true
+		}
+	}
+	if partialSort {
+		for i := 1; i < serversCount; i++ {
+			if serversInfo.inner[i-1].rtt.Value() > serversInfo.inner[i].rtt.Value() {
+				serversInfo.inner[i-1], serversInfo.inner[i] = serversInfo.inner[i], serversInfo.inner[i-1]
+			}
+		}
 	}
 	switch serversInfo.lbStrategy {
 	case LBStrategyFastest:
@@ -315,7 +335,7 @@ func (serversInfo *ServersInfo) fetchDoHServerInfo(proxy *Proxy, name string, st
 
 func (serverInfo *ServerInfo) noticeFailure(proxy *Proxy) {
 	serverInfo.Lock()
-	serverInfo.rtt.Set(float64(proxy.timeout.Nanoseconds()))
+	serverInfo.rtt.Add(float64(proxy.timeout.Nanoseconds() / 1000000))
 	serverInfo.Unlock()
 }
 
@@ -328,9 +348,10 @@ func (serverInfo *ServerInfo) noticeBegin(proxy *Proxy) {
 func (serverInfo *ServerInfo) noticeSuccess(proxy *Proxy) {
 	now := time.Now()
 	serverInfo.Lock()
-	elapsed := now.Sub(serverInfo.lastActionTS) / 1024
-	if elapsed > 0 {
-		serverInfo.rtt.Add(float64(elapsed))
+	elapsed := now.Sub(serverInfo.lastActionTS)
+	elapsedMs := elapsed.Nanoseconds() / 1000000
+	if elapsedMs > 0 && elapsed < proxy.timeout {
+		serverInfo.rtt.Add(float64(elapsedMs))
 	}
 	serverInfo.Unlock()
 }