Merge branch 'lbstrategy-interface'

* lbstrategy-interface:
  Support power-of-<arbitrary number>
  Use an interface for load-balancing strategies

dnscrypt-proxy/config.go

@@ -358,23 +358,30 @@ func ConfigLoad(proxy *Proxy, flags *ConfigFlags) error {
 	if len(config.ListenAddresses) == 0 && len(config.LocalDoH.ListenAddresses) == 0 {
 		dlog.Debug("No local IP/port configured")
 	}
-
+	lbStrategy := LBStrategy(DefaultLBStrategy)
-	lbStrategy := DefaultLBStrategy
+	switch lbStrategyStr := strings.ToLower(config.LBStrategy); lbStrategyStr {
-	switch strings.ToLower(config.LBStrategy) {
 	case "":
 		// default
 	case "p2":
-		lbStrategy = LBStrategyP2
+		lbStrategy = LBStrategyP2{}
 	case "ph":
-		lbStrategy = LBStrategyPH
+		lbStrategy = LBStrategyPH{}
 	case "fastest":
 	case "first":
-		lbStrategy = LBStrategyFirst
+		lbStrategy = LBStrategyFirst{}
 	case "random":
-		lbStrategy = LBStrategyRandom
+		lbStrategy = LBStrategyRandom{}
 	default:
+		if strings.HasPrefix(lbStrategyStr, "p") {
+			n, err := strconv.ParseInt(strings.TrimPrefix(lbStrategyStr, "p"), 10, 32)
+			if err != nil || n <= 0 {
+				dlog.Fatalf("Invalid load balancing strategy: [%s]", config.LBStrategy)
+			}
+			lbStrategy = LBStrategyPN{n: int(n)}
+		} else {
 			dlog.Warnf("Unknown load balancing strategy: [%s]", config.LBStrategy)
 		}
+	}
 	proxy.serversInfo.lbStrategy = lbStrategy
 	proxy.serversInfo.lbEstimator = config.LBEstimator
 

dnscrypt-proxy/serversInfo.go

@@ -62,17 +62,41 @@ type ServerInfo struct {
 	DOHClientCreds     DOHClientCreds
 }
 
-type LBStrategy int
+type LBStrategy interface {
+	getCandidate(serversCount int) int
+}
 
-const (
+type LBStrategyP2 struct{}
-	LBStrategyNone = LBStrategy(iota)
-	LBStrategyP2
-	LBStrategyPH
-	LBStrategyFirst
-	LBStrategyRandom
-)
 
-const DefaultLBStrategy = LBStrategyP2
+func (LBStrategyP2) getCandidate(serversCount int) int {
+	return rand.Intn(Min(serversCount, 2))
+}
+
+type LBStrategyPN struct{ n int }
+
+func (s LBStrategyPN) getCandidate(serversCount int) int {
+	return rand.Intn(Min(serversCount, s.n))
+}
+
+type LBStrategyPH struct{}
+
+func (LBStrategyPH) getCandidate(serversCount int) int {
+	return rand.Intn(Max(Min(serversCount, 2), serversCount/2))
+}
+
+type LBStrategyFirst struct{}
+
+func (LBStrategyFirst) getCandidate(int) int {
+	return 0
+}
+
+type LBStrategyRandom struct{}
+
+func (LBStrategyRandom) getCandidate(serversCount int) int {
+	return rand.Intn(serversCount)
+}
+
+var DefaultLBStrategy = LBStrategyP2{}
 
 type ServersInfo struct {
 	sync.RWMutex
@@ -209,17 +233,7 @@ func (serversInfo *ServersInfo) getOne() *ServerInfo {
 	if serversInfo.lbEstimator {
 		serversInfo.estimatorUpdate()
 	}
-	var candidate int
+	candidate := serversInfo.lbStrategy.getCandidate(serversCount)
-	switch serversInfo.lbStrategy {
-	case LBStrategyFirst:
-		candidate = 0
-	case LBStrategyPH:
-		candidate = rand.Intn(Max(Min(serversCount, 2), serversCount/2))
-	case LBStrategyRandom:
-		candidate = rand.Intn(serversCount)
-	default:
-		candidate = rand.Intn(Min(serversCount, 2))
-	}
 	serverInfo := serversInfo.inner[candidate]
 	dlog.Debugf("Using candidate [%s] RTT: %d", (*serverInfo).Name, int((*serverInfo).rtt.Value()))
 	serversInfo.Unlock()