package roundrobin import ( "fmt" "net/http" "net/url" "sync" "time" "github.com/mailgun/timetools" "github.com/vulcand/oxy/memmetrics" "github.com/vulcand/oxy/utils" ) // RebalancerOption - functional option setter for rebalancer type RebalancerOption func(*Rebalancer) error // Meter measures server peformance and returns it's relative value via rating type Meter interface { Rating() float64 Record(int, time.Duration) IsReady() bool } type NewMeterFn func() (Meter, error) // Rebalancer increases weights on servers that perform better than others. It also rolls back to original weights // if the servers have changed. It is designed as a wrapper on top of the roundrobin. type Rebalancer struct { // mutex mtx *sync.Mutex // As usual, control time in tests clock timetools.TimeProvider // Time that freezes state machine to accumulate stats after updating the weights backoffDuration time.Duration // Timer is set to give probing some time to take place timer time.Time // server records that remember original weights servers []*rbServer // next is internal load balancer next in chain next balancerHandler // errHandler is HTTP handler called in case of errors errHandler utils.ErrorHandler log utils.Logger ratings []float64 // creates new meters newMeter NewMeterFn // sticky session object ss *StickySession } func RebalancerLogger(log utils.Logger) RebalancerOption { return func(r *Rebalancer) error { r.log = log return nil } } func RebalancerClock(clock timetools.TimeProvider) RebalancerOption { return func(r *Rebalancer) error { r.clock = clock return nil } } func RebalancerBackoff(d time.Duration) RebalancerOption { return func(r *Rebalancer) error { r.backoffDuration = d return nil } } func RebalancerMeter(newMeter NewMeterFn) RebalancerOption { return func(r *Rebalancer) error { r.newMeter = newMeter return nil } } // RebalancerErrorHandler is a functional argument that sets error handler of the server func RebalancerErrorHandler(h utils.ErrorHandler) RebalancerOption { return func(r *Rebalancer) error { r.errHandler = h return nil } } func RebalancerStickySession(ss *StickySession) RebalancerOption { return func(r *Rebalancer) error { r.ss = ss return nil } } func NewRebalancer(handler balancerHandler, opts ...RebalancerOption) (*Rebalancer, error) { rb := &Rebalancer{ mtx: &sync.Mutex{}, next: handler, ss: nil, } for _, o := range opts { if err := o(rb); err != nil { return nil, err } } if rb.clock == nil { rb.clock = &timetools.RealTime{} } if rb.backoffDuration == 0 { rb.backoffDuration = 10 * time.Second } if rb.log == nil { rb.log = &utils.NOPLogger{} } if rb.newMeter == nil { rb.newMeter = func() (Meter, error) { rc, err := memmetrics.NewRatioCounter(10, time.Second, memmetrics.RatioClock(rb.clock)) if err != nil { return nil, err } return &codeMeter{ r: rc, codeS: http.StatusInternalServerError, codeE: http.StatusGatewayTimeout + 1, }, nil } } if rb.errHandler == nil { rb.errHandler = utils.DefaultHandler } return rb, nil } func (rb *Rebalancer) Servers() []*url.URL { rb.mtx.Lock() defer rb.mtx.Unlock() return rb.next.Servers() } func (rb *Rebalancer) ServeHTTP(w http.ResponseWriter, req *http.Request) { pw := &utils.ProxyWriter{W: w} start := rb.clock.UtcNow() // make shallow copy of request before changing anything to avoid side effects newReq := *req stuck := false if rb.ss != nil { cookie_url, present, err := rb.ss.GetBackend(&newReq, rb.Servers()) if err != nil { rb.errHandler.ServeHTTP(w, req, err) return } if present { newReq.URL = cookie_url stuck = true } } if !stuck { url, err := rb.next.NextServer() if err != nil { rb.errHandler.ServeHTTP(w, req, err) return } if rb.ss != nil { rb.ss.StickBackend(url, &w) } newReq.URL = url } rb.next.Next().ServeHTTP(pw, &newReq) rb.recordMetrics(newReq.URL, pw.Code, rb.clock.UtcNow().Sub(start)) rb.adjustWeights() } func (rb *Rebalancer) recordMetrics(u *url.URL, code int, latency time.Duration) { rb.mtx.Lock() defer rb.mtx.Unlock() if srv, i := rb.findServer(u); i != -1 { srv.meter.Record(code, latency) } } func (rb *Rebalancer) reset() { for _, s := range rb.servers { s.curWeight = s.origWeight rb.next.UpsertServer(s.url, Weight(s.origWeight)) } rb.timer = rb.clock.UtcNow().Add(-1 * time.Second) rb.ratings = make([]float64, len(rb.servers)) } func (rb *Rebalancer) Wrap(next balancerHandler) error { if rb.next != nil { return fmt.Errorf("already bound to %T", rb.next) } rb.next = next return nil } func (rb *Rebalancer) UpsertServer(u *url.URL, options ...ServerOption) error { rb.mtx.Lock() defer rb.mtx.Unlock() if err := rb.next.UpsertServer(u, options...); err != nil { return err } weight, _ := rb.next.ServerWeight(u) if err := rb.upsertServer(u, weight); err != nil { rb.next.RemoveServer(u) return err } rb.reset() return nil } func (rb *Rebalancer) RemoveServer(u *url.URL) error { rb.mtx.Lock() defer rb.mtx.Unlock() return rb.removeServer(u) } func (rb *Rebalancer) removeServer(u *url.URL) error { _, i := rb.findServer(u) if i == -1 { return fmt.Errorf("%v not found", u) } if err := rb.next.RemoveServer(u); err != nil { return err } rb.servers = append(rb.servers[:i], rb.servers[i+1:]...) rb.reset() return nil } func (rb *Rebalancer) upsertServer(u *url.URL, weight int) error { if s, i := rb.findServer(u); i != -1 { s.origWeight = weight } meter, err := rb.newMeter() if err != nil { return err } rbSrv := &rbServer{ url: utils.CopyURL(u), origWeight: weight, curWeight: weight, meter: meter, } rb.servers = append(rb.servers, rbSrv) return nil } func (r *Rebalancer) findServer(u *url.URL) (*rbServer, int) { if len(r.servers) == 0 { return nil, -1 } for i, s := range r.servers { if sameURL(u, s.url) { return s, i } } return nil, -1 } // Called on every load balancer ServeHTTP call, returns the suggested weights // on every call, can adjust weights if needed. func (rb *Rebalancer) adjustWeights() { rb.mtx.Lock() defer rb.mtx.Unlock() // In this case adjusting weights would have no effect, so do nothing if len(rb.servers) < 2 { return } // Metrics are not ready if !rb.metricsReady() { return } if !rb.timerExpired() { return } if rb.markServers() { if rb.setMarkedWeights() { rb.setTimer() } } else { // No servers that are different by their quality, so converge weights if rb.convergeWeights() { rb.setTimer() } } } func (rb *Rebalancer) applyWeights() { for _, srv := range rb.servers { rb.log.Infof("upsert server %v, weight %v", srv.url, srv.curWeight) rb.next.UpsertServer(srv.url, Weight(srv.curWeight)) } } func (rb *Rebalancer) setMarkedWeights() bool { changed := false // Increase weights on servers marked as good for _, srv := range rb.servers { if srv.good { weight := increase(srv.curWeight) if weight <= FSMMaxWeight { rb.log.Infof("increasing weight of %v from %v to %v", srv.url, srv.curWeight, weight) srv.curWeight = weight changed = true } } } if changed { rb.normalizeWeights() rb.applyWeights() return true } return false } func (rb *Rebalancer) setTimer() { rb.timer = rb.clock.UtcNow().Add(rb.backoffDuration) } func (rb *Rebalancer) timerExpired() bool { return rb.timer.Before(rb.clock.UtcNow()) } func (rb *Rebalancer) metricsReady() bool { for _, s := range rb.servers { if !s.meter.IsReady() { return false } } return true } // markServers splits servers into two groups of servers with bad and good failure rate. // It does compare relative performances of the servers though, so if all servers have approximately the same error rate // this function returns the result as if all servers are equally good. func (rb *Rebalancer) markServers() bool { for i, srv := range rb.servers { rb.ratings[i] = srv.meter.Rating() } g, b := memmetrics.SplitFloat64(splitThreshold, 0, rb.ratings) for i, srv := range rb.servers { if g[rb.ratings[i]] { srv.good = true } else { srv.good = false } } if len(g) != 0 && len(b) != 0 { rb.log.Infof("bad: %v good: %v, ratings: %v", b, g, rb.ratings) } return len(g) != 0 && len(b) != 0 } func (rb *Rebalancer) convergeWeights() bool { // If we have previoulsy changed servers try to restore weights to the original state changed := false for _, s := range rb.servers { if s.origWeight == s.curWeight { continue } changed = true newWeight := decrease(s.origWeight, s.curWeight) rb.log.Infof("decreasing weight of %v from %v to %v", s.url, s.curWeight, newWeight) s.curWeight = newWeight } if !changed { return false } rb.normalizeWeights() rb.applyWeights() return true } func (rb *Rebalancer) weightsGcd() int { divisor := -1 for _, w := range rb.servers { if divisor == -1 { divisor = w.curWeight } else { divisor = gcd(divisor, w.curWeight) } } return divisor } func (rb *Rebalancer) normalizeWeights() { gcd := rb.weightsGcd() if gcd <= 1 { return } for _, s := range rb.servers { s.curWeight = s.curWeight / gcd } } func increase(weight int) int { return weight * FSMGrowFactor } func decrease(target, current int) int { adjusted := current / FSMGrowFactor if adjusted < target { return target } else { return adjusted } } // rebalancer server record that keeps track of the original weight supplied by user type rbServer struct { url *url.URL origWeight int // original weight supplied by user curWeight int // current weight good bool meter Meter } const ( // This is the maximum weight that handler will set for the server FSMMaxWeight = 4096 // Multiplier for the server weight FSMGrowFactor = 4 ) type codeMeter struct { r *memmetrics.RatioCounter codeS int codeE int } func (n *codeMeter) Rating() float64 { return n.r.Ratio() } func (n *codeMeter) Record(code int, d time.Duration) { if code >= n.codeS && code < n.codeE { n.r.IncA(1) } else { n.r.IncB(1) } } func (n *codeMeter) IsReady() bool { return n.r.IsReady() } // splitThreshold tells how far the value should go from the median + median absolute deviation before it is considered an outlier const splitThreshold = 1.5