Add backend weight round robin select (#34)

* Add upstream selector, there are two selector now: - random selector - weight random selector random selector will choose upstream at random; weight random selector will choose upstream at random with weight Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Rewrite config and config file example, prepare for weight round robbin selector Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Replace bad implement of weight random selector with weight round robbin selector, the algorithm is nginx weight round robbin like Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Use new config module Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Disable deprecated DualStack set Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Fix typo Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Optimize upstreamSelector judge Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Fix typo Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Add config timeout unit tips Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Set wrr http client timeout to replace http request timeout Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Add weight value range Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Add a line ending for .gitignore Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Optimize config file style Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Modify Weight type to int32 Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Add upstreamError Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Rewrite Selector interface and wrr implement Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Use http module predefined constant to judge req.response.StatusCode Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Use Selector.ReportUpstreamError to report upstream error for evaluation loop in real time Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Make client selector field private Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Replace config file url to URL Add miss space for 'weight= 50' Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Rewrite Selector.ReportUpstreamError to Selector.ReportUpstreamStatus, report upstream ok in real time Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Fix checkIETFResponse: if upstream OK, won't increase weight Signed-off-by: Sherlock Holo <sherlockya@gmail.com> * Fix typo Signed-off-by: Sherlock Holo <sherlockya@gmail.com>
2026-05-15 01:35:41 +00:00 · 2019-03-09 18:12:45 +08:00
parent 8f2004d1de
commit fec1e84d5e
14 changed files with 552 additions and 122 deletions
@@ -0,0 +1,50 @@
+package selector
+
+import (
+	"errors"
+	"math/rand"
+	"time"
+)
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+type RandomSelector struct {
+	upstreams []*Upstream
+}
+
+func NewRandomSelector() *RandomSelector {
+	return new(RandomSelector)
+}
+
+func (rs *RandomSelector) Add(url string, upstreamType UpstreamType) (err error) {
+	switch upstreamType {
+	case Google:
+		rs.upstreams = append(rs.upstreams, &Upstream{
+			Type:        Google,
+			Url:         url,
+			RequestType: "application/dns-json",
+		})
+
+	case IETF:
+		rs.upstreams = append(rs.upstreams, &Upstream{
+			Type:        IETF,
+			Url:         url,
+			RequestType: "application/dns-message",
+		})
+
+	default:
+		return errors.New("unknown upstream type")
+	}
+
+	return nil
+}
+
+func (rs *RandomSelector) Get() *Upstream {
+	return rs.upstreams[rand.Intn(len(rs.upstreams)-1)]
+}
+
+func (rs *RandomSelector) StartEvaluate() {}
+
+func (rs *RandomSelector) ReportUpstreamStatus(upstream *Upstream, upstreamStatus upstreamStatus) {}
@@ -0,0 +1,12 @@
+package selector
+
+type Selector interface {
+	// Get returns a upstream
+	Get() *Upstream
+
+	// StartEvaluate start upstream evaluation loop
+	StartEvaluate()
+
+	// ReportUpstreamStatus report upstream status
+	ReportUpstreamStatus(upstream *Upstream, upstreamStatus upstreamStatus)
+}
@@ -0,0 +1,28 @@
+package selector
+
+import "fmt"
+
+type UpstreamType int
+
+const (
+	Google UpstreamType = iota
+	IETF
+)
+
+var typeMap = map[UpstreamType]string{
+	Google: "Google",
+	IETF:   "IETF",
+}
+
+type Upstream struct {
+	Type            UpstreamType
+	Url             string
+	RequestType     string
+	weight          int32
+	effectiveWeight int32
+	currentWeight   int32
+}
+
+func (u Upstream) String() string {
+	return fmt.Sprintf("upstream type: %s, upstream url: %s", typeMap[u.Type], u.Url)
+}
@@ -0,0 +1,14 @@
+package selector
+
+type upstreamStatus int
+
+const (
+	// when query upstream timeout, usually upstream is unavailable for a long time
+	Timeout upstreamStatus = iota
+
+	// when query upstream return 5xx response, upstream still alive, maybe just a lof of query for him
+	Error
+
+	// when query upstream ok, means upstream is available
+	OK
+)
@@ -0,0 +1,192 @@
+package selector
+
+import (
+	"encoding/json"
+	"errors"
+	"net/http"
+	"sync/atomic"
+	"time"
+)
+
+type WeightRoundRobinSelector struct {
+	upstreams []*Upstream // upstreamsInfo
+	client    http.Client // http client to check the upstream
+}
+
+func NewWeightRoundRobinSelector(timeout time.Duration) *WeightRoundRobinSelector {
+	return &WeightRoundRobinSelector{
+		client: http.Client{Timeout: timeout},
+	}
+}
+
+func (ws *WeightRoundRobinSelector) Add(url string, upstreamType UpstreamType, weight int32) (err error) {
+	switch upstreamType {
+	case Google:
+		ws.upstreams = append(ws.upstreams, &Upstream{
+			Type:            Google,
+			Url:             url,
+			RequestType:     "application/dns-json",
+			weight:          weight,
+			effectiveWeight: weight,
+		})
+
+	case IETF:
+		ws.upstreams = append(ws.upstreams, &Upstream{
+			Type:            IETF,
+			Url:             url,
+			RequestType:     "application/dns-message",
+			weight:          weight,
+			effectiveWeight: weight,
+		})
+
+	default:
+		return errors.New("unknown upstream type")
+	}
+
+	return nil
+}
+
+// COW, avoid concurrent read write upstreams
+func (ws *WeightRoundRobinSelector) StartEvaluate() {
+	go func() {
+		for {
+			for i := range ws.upstreams {
+				upstreamUrl := ws.upstreams[i].Url
+				var acceptType string
+
+				switch ws.upstreams[i].Type {
+				case Google:
+					upstreamUrl += "?name=www.example.com&type=A"
+					acceptType = "application/dns-json"
+
+				case IETF:
+					// www.example.com
+					upstreamUrl += "?dns=q80BAAABAAAAAAAAA3d3dwdleGFtcGxlA2NvbQAAAQAB"
+					acceptType = "application/dns-message"
+				}
+
+				req, err := http.NewRequest(http.MethodGet, upstreamUrl, nil)
+				if err != nil {
+					/*log.Println("upstream:", upstreamUrl, "type:", typeMap[upstream.Type], "check failed:", err)
+					continue*/
+
+					// should I only log it? But if there is an error, I think when query the server will return error too
+					panic("upstream: " + upstreamUrl + " type: " + typeMap[ws.upstreams[i].Type] + " check failed: " + err.Error())
+				}
+
+				req.Header.Set("accept", acceptType)
+
+				resp, err := ws.client.Do(req)
+				if err != nil {
+					// should I check error in detail?
+					if atomic.AddInt32(&ws.upstreams[i].effectiveWeight, -10) < 0 {
+						atomic.StoreInt32(&ws.upstreams[i].effectiveWeight, 0)
+					}
+					continue
+				}
+
+				switch ws.upstreams[i].Type {
+				case Google:
+					checkGoogleResponse(resp, ws.upstreams[i])
+
+				case IETF:
+					checkIETFResponse(resp, ws.upstreams[i])
+				}
+			}
+
+			time.Sleep(30 * time.Second)
+		}
+	}()
+}
+
+// nginx wrr like
+func (ws *WeightRoundRobinSelector) Get() *Upstream {
+	var (
+		total             int32
+		bestUpstreamIndex = -1
+	)
+
+	for i := range ws.upstreams {
+		effectiveWeight := atomic.LoadInt32(&ws.upstreams[i].effectiveWeight)
+		ws.upstreams[i].currentWeight += effectiveWeight
+		total += effectiveWeight
+
+		if bestUpstreamIndex == -1 || ws.upstreams[i].currentWeight > ws.upstreams[bestUpstreamIndex].currentWeight {
+			bestUpstreamIndex = i
+		}
+	}
+
+	ws.upstreams[bestUpstreamIndex].currentWeight -= total
+
+	return ws.upstreams[bestUpstreamIndex]
+}
+
+func (ws *WeightRoundRobinSelector) ReportUpstreamStatus(upstream *Upstream, upstreamStatus upstreamStatus) {
+	switch upstreamStatus {
+	case Timeout:
+		if atomic.AddInt32(&upstream.effectiveWeight, -10) < 0 {
+			atomic.StoreInt32(&upstream.effectiveWeight, 0)
+		}
+
+	case Error:
+		if atomic.AddInt32(&upstream.effectiveWeight, -5) < 0 {
+			atomic.StoreInt32(&upstream.effectiveWeight, 0)
+		}
+
+	case OK:
+		if atomic.AddInt32(&upstream.effectiveWeight, 2) > upstream.weight {
+			atomic.StoreInt32(&upstream.effectiveWeight, upstream.weight)
+		}
+	}
+}
+
+func checkGoogleResponse(resp *http.Response, upstream *Upstream) {
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		// server error
+		if atomic.AddInt32(&upstream.effectiveWeight, -5) < 0 {
+			atomic.StoreInt32(&upstream.effectiveWeight, 0)
+		}
+		return
+	}
+
+	m := make(map[string]interface{})
+	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
+		// should I check error in detail?
+		if atomic.AddInt32(&upstream.effectiveWeight, -1) < 0 {
+			atomic.StoreInt32(&upstream.effectiveWeight, 0)
+		}
+		return
+	}
+
+	if status, ok := m["status"]; ok {
+		if statusNum, ok := status.(int); ok && statusNum == 0 {
+			if atomic.AddInt32(&upstream.effectiveWeight, 5) > upstream.weight {
+				atomic.StoreInt32(&upstream.effectiveWeight, upstream.weight)
+			}
+			return
+		}
+	}
+
+	// should I check error in detail?
+	if atomic.AddInt32(&upstream.effectiveWeight, -1) < 0 {
+		atomic.StoreInt32(&upstream.effectiveWeight, 0)
+	}
+}
+
+func checkIETFResponse(resp *http.Response, upstream *Upstream) {
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		// server error
+		if atomic.AddInt32(&upstream.effectiveWeight, -5) < 0 {
+			atomic.StoreInt32(&upstream.effectiveWeight, 0)
+		}
+		return
+	}
+
+	if atomic.AddInt32(&upstream.effectiveWeight, 5) > upstream.weight {
+		atomic.StoreInt32(&upstream.effectiveWeight, upstream.weight)
+	}
+}