Merge pull request #120 from domgoer/master

fix: use unsupported loadbalance will cause nil pointer

docs/server.md

@@ -12,6 +12,9 @@ Server地址，格式为："IP:PORT"。
 ## Protocol
 Server的接口协议，目前支持HTTP。
 
+## Weight
+Weight 服务器的权重（当该服务器所属的集群负载方式是权重轮询时则需要配置）
+
 ## MaxQPS
 Server能够支持的最大QPS，用于流控。Gateway采用令牌桶算法，根据QPS限制流量，保护后端Server被压垮。
 

pkg/client/server.go

@@ -80,6 +80,12 @@ func (sb *ServerBuilder) MaxQPS(max int64) *ServerBuilder {
 	return sb
 }
 
+// Weight set robin weight
+func (sb *ServerBuilder) Weight(weight int64) *ServerBuilder {
+	sb.value.Weight = weight
+	return sb
+}
+
 // NoCircuitBreaker no circuit breaker
 func (sb *ServerBuilder) NoCircuitBreaker() *ServerBuilder {
 	sb.value.CircuitBreaker = nil

pkg/lb/lb.go

@@ -15,20 +15,25 @@ var (
 	// LBS map loadBalance name and process function
 	LBS = map[metapb.LoadBalance]func() LoadBalance{
 		metapb.RoundRobin: NewRoundRobin,
+		metapb.WightRobin: NewWeightRobin,
 	}
 )
 
-// LoadBalance loadBalance interface
+// LoadBalance loadBalance interface returns selected server's id
 type LoadBalance interface {
-	Select(req *fasthttp.Request, servers *list.List) int
+	Select(req *fasthttp.Request, servers *list.List) uint64
 }
 
 // GetSupportLBS return supported loadBalances
 func GetSupportLBS() []metapb.LoadBalance {
 	return supportLbs
 }
 
-// NewLoadBalance create a LoadBalance
+// NewLoadBalance create a LoadBalance,if LoadBalance function is not supported
+// it will return NewRoundRobin
 func NewLoadBalance(name metapb.LoadBalance) LoadBalance {
-	return LBS[name]()
+	if l, ok := LBS[name]; ok {
+		return l()
+	}
+	return NewRoundRobin()
 }

pkg/lb/roundrobin.go

@@ -2,6 +2,8 @@ package lb
 
 import (
 	"container/list"
+	"github.com/fagongzi/gateway/pkg/pb/metapb"
+	"github.com/fagongzi/util/collection"
 	"sync/atomic"
 
 	"github.com/valyala/fasthttp"
@@ -23,12 +25,19 @@ func NewRoundRobin() LoadBalance {
 }
 
 // Select select a server from servers using RoundRobin
-func (rr RoundRobin) Select(req *fasthttp.Request, servers *list.List) int {
+func (rr RoundRobin) Select(req *fasthttp.Request, servers *list.List) uint64 {
 	l := uint64(servers.Len())
 
 	if 0 >= l {
-		return -1
+		return 0
 	}
 
-	return int(atomic.AddUint64(rr.ops, 1) % l)
+	idx := int(atomic.AddUint64(rr.ops, 1) % l)
+
+	v := collection.Get(servers, idx).Value
+	if v == nil {
+		return 0
+	}
+
+	return v.(*metapb.Server).ID
 }

pkg/lb/weightrobin.go

@@ -0,0 +1,61 @@
+package lb
+
+import (
+	"container/list"
+	"github.com/fagongzi/gateway/pkg/pb/metapb"
+	"github.com/valyala/fasthttp"
+)
+
+// WeightRobin weight robin loadBalance impl
+type WeightRobin struct {
+	opts map[uint64]*weightRobin
+}
+
+// weightRobin used to save the weight info of server
+type weightRobin struct {
+	effectiveWeight int64
+	currentWeight   int64
+}
+
+// NewWeightRobin create a WeightRobin
+func NewWeightRobin() LoadBalance {
+	return &WeightRobin{
+		opts: make(map[uint64]*weightRobin, 1024),
+	}
+}
+
+// Select select a server from servers using WeightRobin
+func (w *WeightRobin) Select(req *fasthttp.Request, servers *list.List) (best uint64) {
+	var total int64
+
+	for iter := servers.Back(); iter != nil; iter = iter.Prev() {
+		svr := iter.Value.(*metapb.Server)
+
+		id := svr.ID
+		if _, ok := w.opts[id]; !ok {
+			w.opts[id] = &weightRobin{
+				effectiveWeight: svr.Weight,
+			}
+		}
+
+		wt := w.opts[id]
+		wt.currentWeight += wt.effectiveWeight
+		total += wt.effectiveWeight
+
+		if wt.effectiveWeight < svr.Weight {
+			wt.effectiveWeight++
+		}
+
+		if best == 0 || w.opts[uint64(best)] == nil || wt.currentWeight > w.opts[best].currentWeight {
+			best = id
+		}
+	}
+
+	if best == 0 {
+		return 0
+	}
+
+	w.opts[best].currentWeight -= total
+
+	return best
+}

pkg/lb/weightrobin_test.go

@@ -0,0 +1,69 @@
+package lb
+
+import (
+	"container/list"
+	"github.com/fagongzi/gateway/pkg/pb/metapb"
+	"github.com/valyala/fasthttp"
+	"testing"
+)
+
+func TestWeightRobin_Select(t *testing.T) {
+	li := list.New()
+
+	li.PushBack(&metapb.Server{
+		ID:     1,
+		Weight: 20,
+	})
+	li.PushBack(&metapb.Server{
+		ID:     2,
+		Weight: 10,
+	})
+	li.PushBack(&metapb.Server{
+		ID:     3,
+		Weight: 35,
+	})
+	li.PushBack(&metapb.Server{
+		ID:     4,
+		Weight: 5,
+	})
+
+	type fields struct {
+		opts map[uint64]*weightRobin
+	}
+	type args struct {
+		req     *fasthttp.Request
+		servers *list.List
+	}
+	tests := []struct {
+		name     string
+		fields   fields
+		args     args
+		wantBest []int
+	}{
+		{
+			name:   "test_case_1",
+			fields: struct{ opts map[uint64]*weightRobin }{opts: make(map[uint64]*weightRobin, 50)},
+			args: struct {
+				req     *fasthttp.Request
+				servers *list.List
+			}{req: nil, servers: li},
+			wantBest: []int{20, 10, 35, 5},
+		},
+	}
+	for _, tt := range tests {
+		var res = make(map[uint64]int)
+		t.Run(tt.name, func(t *testing.T) {
+			w := &WeightRobin{
+				opts: tt.fields.opts,
+			}
+			for i := 0; i < 70; i++ {
+				res[w.Select(tt.args.req, tt.args.servers)]++
+			}
+		})
+		for k, v := range res {
+			if tt.wantBest[k-1] != v {
+				t.Errorf("WeightRobin.Select() = %v, want %v", res, tt.wantBest)
+			}
+		}
+	}
+}