diff --git a/rule/forward.go b/rule/forward.go index 2304220..00f7956 100644 --- a/rule/forward.go +++ b/rule/forward.go @@ -27,6 +27,7 @@ type Forwarder struct { latency int64 intface string // local interface or ip address handlers []StatusHandler + weight int64 } // ForwarderFromURL parses `forward=` command value and returns a new forwarder. @@ -97,6 +98,17 @@ func (f *Forwarder) parseOption(option string) error { } f.SetPriority(uint32(priority)) + var weight int64 + w := query.Get("weight") + if w != "" { + weight, err = strconv.ParseInt(w, 10, 32) + if err == nil && weight > 0 { + f.SetWeight(weight) + } else if weight < 0 { + log.F("weight should be more than 0, ignore weight value") + } + } + f.intface = query.Get("interface") return err @@ -204,3 +216,11 @@ func (f *Forwarder) Latency() int64 { func (f *Forwarder) SetLatency(l int64) { atomic.StoreInt64(&f.latency, l) } +func (f *Forwarder) Weight() int64 { + return f.weight +} + +// SetWeight sets the weight of forwarder. +func (f *Forwarder) SetWeight(w int64) { + f.weight = w +} diff --git a/rule/group.go b/rule/group.go index fc9607b..943c2ff 100644 --- a/rule/group.go +++ b/rule/group.go @@ -3,6 +3,7 @@ package rule import ( "errors" "hash/fnv" + "math/rand" "net" "net/url" "path/filepath" @@ -25,14 +26,15 @@ func (p priSlice) Swap(i, j int) { p[i], p[j] = p[j], p[i] } // FwdrGroup is a forwarder group. type FwdrGroup struct { - name string - config *Strategy - fwdrs priSlice - avail []*Forwarder // available forwarders - mu sync.RWMutex - index uint32 - priority uint32 - next func(addr string) *Forwarder + name string + config *Strategy + fwdrs priSlice + avail []*Forwarder // available forwarders + mu sync.RWMutex + index uint32 + priority uint32 + next func(addr string) *Forwarder + randomGenerator *rand.Rand } // NewFwdrGroup returns a new forward group. @@ -88,6 +90,10 @@ func newFwdrGroup(name string, fwdrs []*Forwarder, c *Strategy) *FwdrGroup { case "dh": p.next = p.scheduleDH log.F("[strategy] %s: %d forwarders forward in destination hashing mode.", name, count) + case "wb": + p.next = p.scheduleWB + p.randomGenerator = rand.New(rand.NewSource(time.Now().UnixNano())) + log.F("[strategy] %s: %d forwarders forward in Weight base mode.", name, count) default: p.next = p.scheduleRR log.F("[strategy] %s: not supported forward mode '%s', use round robin mode for %d forwarders.", name, c.Strategy, count) @@ -316,3 +322,30 @@ func (p *FwdrGroup) scheduleDH(dstAddr string) *Forwarder { fnv1a.Write([]byte(dstAddr)) return p.avail[fnv1a.Sum32()%uint32(len(p.avail))] } + +// Weight base +func (p *FwdrGroup) scheduleWB(dstAddr string) *Forwarder { + totalWeight := int64(0) + for _, f := range p.avail { + totalWeight += f.Weight() + } + if totalWeight <= 0 { + log.F("total weight is zero switch to rr mode") + return p.scheduleRR(dstAddr) + } + r := p.randomGenerator.Int63n(totalWeight) + var sum int64 = 0 + for _, f := range p.avail { + weight := f.Weight() + if weight == 0 { + return f + } + sum += weight + if r < sum { + return f + } + } + // Shouldn't reach here + //just in case + return p.avail[0] +}