<commit_before>package roundrobin

import (
	"fmt"
	log "github.com/mailgun/gotools-log"
	timetools "github.com/mailgun/gotools-time"
	"math"
	"time"
)

type FSMState int

const (
	FSMStart    = iota // initial state of the fsm
	FSMProbing  = iota // fsm is trying some theory
	FSMRollback = iota // fsm is rolling back
	FSMRevert   = iota // fsm is getting back to the original state
)

const (
	FSMMaxWeight = 16384
)

// This is the tiny DFA that tries to play with weights to improve over the overall error rate
// to see if it helps and falls back if taking the load off the bad upstream makes the situation worse.
type FSMHandler struct {
	timeProvider        timetools.TimeProvider
	backoffDuration     time.Duration      // duration that we use to backoff or apply any theory
	state               FSMState           // Current state of the state machione
	timer               time.Time          // Timer is set to give probing some time to take place
	probedGoodEndpoints []*changedEndpoint // Probing changes endpoint weights and remembers the weight so it can go back in case of failure
	weightChanges       int
}

type changedEndpoint struct {
	failRatioBefore float64
	endpoint        *WeightedEndpoint
	weightBefore    int
}

func NewFSMHandler() (*FSMHandler, error) {
	return NewFSMHandlerWithOptions(&timetools.RealTime{}, 10*time.Second)
}

func NewFSMHandlerWithOptions(timeProvider timetools.TimeProvider, duration time.Duration) (*FSMHandler, error) {
	if timeProvider == nil {
		return nil, fmt.Errorf("time provider can not be nil")
	}
	if duration < time.Second {
		return nil, fmt.Errorf("supply some backoff duration >= time.Second")
	}
	return &FSMHandler{
		timeProvider:    timeProvider,
		backoffDuration: duration,
	}, nil
}

func (fsm *FSMHandler) reset() {
	fsm.state = FSMStart
	fsm.timer = fsm.timeProvider.UtcNow().Add(-1 * time.Second)
	fsm.probedGoodEndpoints = nil
	fsm.weightChanges = 0
}

func (fsm *FSMHandler) updateWeights(endpoints []*WeightedEndpoint) error {
	if len(endpoints) == 0 {
		fmt.Errorf("No endpoints supplied")
	}
	switch fsm.state {
	case FSMStart:
		return fsm.onStart(endpoints)
	case FSMProbing:
		return fsm.onProbing(endpoints)
	case FSMRollback:
		return fsm.onRollback(endpoints)
	case FSMRevert:
		return fsm.onRollback(endpoints)
	}
	return fmt.Errorf("Invalid state I am in")
}

func (fsm *FSMHandler) onStart(endpoints []*WeightedEndpoint) error {
	failRate := avgFailRate(endpoints)
	// No errors, so let's see if we can recover weights of previosly changed endpoints to the original state
	if failRate == 0 {
		// If we have previoulsy changed endpoints try to restore weights to the original state
		for _, e := range endpoints {
			if e.effectiveWeight != e.weight {
				// Adjust effective weight back to the original weight in stages
				e.setEffectiveWeight(decrease(e.GetOriginalWeight(), e.GetEffectiveWeight()))
				log.Infof("FSMHandler RESTORE %s", e)
				fsm.setTimer()
				fsm.state = FSMRevert
			}
		}
		return nil
	} else {
		log.Infof("FSMHanlder reports average fail rate %f", failRate)
		// Select endpoints with highest error rates and lower their weight
		good, bad := splitEndpoints(endpoints)
		log.Infof("FSMHandler good endpoints: %s", good)
		log.Infof("FSMHandler bad endpoints: %s", bad)
		// No endpoints that are different by their quality
		if len(bad) == 0 || len(good) == 0 {
			return nil
		}
		fsm.probedGoodEndpoints = adjustWeights(good, bad)
		fsm.setTimer()
		fsm.state = FSMProbing
		return nil
	}
}

func (fsm *FSMHandler) onProbing(endpoints []*WeightedEndpoint) error {
	if !fsm.timerExpired() {
		return nil
	}
	// Now revise the good endpoints and see if we made situation worse
	for _, e := range fsm.probedGoodEndpoints {
		if e.failRatioBefore > e.endpoint.failRateMeter.GetRate() {
			// Oops, we made it worse, revert the weights back and go to rollback state
			for _, e := range fsm.probedGoodEndpoints {
				e.endpoint.setEffectiveWeight(e.weightBefore)
			}
			fsm.probedGoodEndpoints = nil
			fsm.state = FSMRollback
			fsm.setTimer()
			return nil
		}
	}
	log.Infof("FSMHandler probing successfull, COMMITING the new rates")
	// We have not made the situation worse, so
	// go back to the starting point and continue the cycle
	fsm.state = FSMStart
	return nil
}

func (fsm *FSMHandler) onRollback(endpoints []*WeightedEndpoint) error {
	if !fsm.timerExpired() {
		return nil
	}
	fsm.state = FSMStart
	return nil
}

func (fsm *FSMHandler) setTimer() {
	fsm.timer = fsm.timeProvider.UtcNow().Add(fsm.backoffDuration)
}

func (fsm *FSMHandler) timerExpired() bool {
	return fsm.timer.Before(fsm.timeProvider.UtcNow())
}

// Splits endpoint into two groups of endpoints with bad performance and good performance. It does compare relative
// performances of the endpoints though, so if all endpoints have the same performance,
func splitEndpoints(endpoints []*WeightedEndpoint) (good []*WeightedEndpoint, bad []*WeightedEndpoint) {
	avg := avgFailRate(endpoints)
	for _, e := range endpoints {
		if greater(e.failRateMeter.GetRate(), avg) {
			bad = append(bad, e)
		} else {
			good = append(good, e)
		}
	}
	return good, bad
}

func adjustWeights(good, bad []*WeightedEndpoint) []*changedEndpoint {
	changedEndpoints := make([]*changedEndpoint, len(good))
	for i, e := range good {
		changed := &changedEndpoint{
			weightBefore:    e.GetEffectiveWeight(),
			failRatioBefore: e.failRateMeter.GetRate(),
			endpoint:        e,
		}
		changedEndpoints[i] = changed
		if increase(e.GetEffectiveWeight()) < FSMMaxWeight {
			e.setEffectiveWeight(increase(e.GetEffectiveWeight()))
			log.Infof("FSMHandler updated weight %s", e)
		}
	}
	return changedEndpoints
}

// Compare two fail rates by neglecting the insignificant differences
func greater(a, b float64) bool {
	return math.Floor(a*100) > math.Floor(b*100)
}

func avgFailRate(endpoints []*WeightedEndpoint) float64 {
	r := float64(0)<commit_msg>Tune the algorithms a bit<commit_after>package roundrobin

import (
	"fmt"
	log "github.com/mailgun/gotools-log"
	timetools "github.com/mailgun/gotools-time"
	"math"
	"time"
)

type FSMState int

const (
	FSMStart    = iota // initial state of the fsm
	FSMProbing  = iota // fsm is trying some theory
	FSMRollback = iota // fsm is rolling back
	FSMRevert   = iota // fsm is getting back to the original state
)

const (
	FSMMaxWeight = 16384
)

// This is the tiny DFA that tries to play with weights to improve over the overall error rate
// to see if it helps and falls back if taking the load off the bad upstream makes the situation worse.
type FSMHandler struct {
	timeProvider        timetools.TimeProvider
	backoffDuration     time.Duration      // duration that we use to backoff or apply any theory
	state               FSMState           // Current state of the state machione
	timer               time.Time          // Timer is set to give probing some time to take place
	probedGoodEndpoints []*changedEndpoint // Probing changes endpoint weights and remembers the weight so it can go back in case of failure
	weightChanges       int
}

type changedEndpoint struct {
	failRatioBefore float64
	endpoint        *WeightedEndpoint
	weightBefore    int
}

func NewFSMHandler() (*FSMHandler, error) {
	return NewFSMHandlerWithOptions(&timetools.RealTime{}, 3*time.Second)
}

func NewFSMHandlerWithOptions(timeProvider timetools.TimeProvider, duration time.Duration) (*FSMHandler, error) {
	if timeProvider == nil {
		return nil, fmt.Errorf("time provider can not be nil")
	}
	if duration < time.Second {
		return nil, fmt.Errorf("supply some backoff duration >= time.Second")
	}
	return &FSMHandler{
		timeProvider:    timeProvider,
		backoffDuration: duration,
	}, nil
}

func (fsm *FSMHandler) reset() {
	fsm.state = FSMStart
	fsm.timer = fsm.timeProvider.UtcNow().Add(-1 * time.Second)
	fsm.probedGoodEndpoints = nil
	fsm.weightChanges = 0
}

func (fsm *FSMHandler) updateWeights(endpoints []*WeightedEndpoint) error {
	if len(endpoints) == 0 {
		fmt.Errorf("No endpoints supplied")
	}
	switch fsm.state {
	case FSMStart:
		return fsm.onStart(endpoints)
	case FSMProbing:
		return fsm.onProbing(endpoints)
	case FSMRollback:
		return fsm.onRollback(endpoints)
	case FSMRevert:
		return fsm.onRollback(endpoints)
	}
	return fmt.Errorf("Invalid state I am in")
}

func (fsm *FSMHandler) onStart(endpoints []*WeightedEndpoint) error {
	failRate := avgFailRate(endpoints)
	// No errors, so let's see if we can recover weights of previosly changed endpoints to the original state
	if failRate == 0 {
		// If we have previoulsy changed endpoints try to restore weights to the original state
		for _, e := range endpoints {
			if e.effectiveWeight != e.weight {
				// Adjust effective weight back to the original weight in stages
				e.setEffectiveWeight(decrease(e.GetOriginalWeight(), e.GetEffectiveWeight()))
				log.Infof("FSMHandler RESTORE %s", e)
				fsm.setTimer()
				fsm.state = FSMRevert
			}
		}
		return nil
	} else {
		log.Infof("FSMHanlder reports average fail rate %f", failRate)
		if !metricsReady(endpoints) {
			log.Infof("FSMHanlder skip cycle, metrics are not ready yet")
			return nil
		}
		// Select endpoints with highest error rates and lower their weight
		good, bad := splitEndpoints(endpoints)
		log.Infof("FSMHandler good endpoints: %s", good)
		log.Infof("FSMHandler bad endpoints: %s", bad)
		// No endpoints that are different by their quality
		if len(bad) == 0 || len(good) == 0 {
			return nil
		}
		fsm.probedGoodEndpoints = adjustWeights(good, bad)
		fsm.setTimer()
		fsm.state = FSMProbing
		return nil
	}
}

func (fsm *FSMHandler) onProbing(endpoints []*WeightedEndpoint) error {
	if !fsm.timerExpired() {
		return nil
	}
	// Now revise the good endpoints and see if we made situation worse
	for _, e := range fsm.probedGoodEndpoints {
		if e.failRatioBefore > e.endpoint.failRateMeter.GetRate() {
			// Oops, we made it worse, revert the weights back and go to rollback state
			for _, e := range fsm.probedGoodEndpoints {
				e.endpoint.setEffectiveWeight(e.weightBefore)
			}
			fsm.probedGoodEndpoints = nil
			fsm.state = FSMRollback
			fsm.setTimer()
			return nil
		}
	}
	log.Infof("FSMHandler probing successfull, COMMITING the new rates")
	// We have not made the situation worse, so
	// go back to the starting point and continue the cycle
	fsm.state = FSMStart
	return nil
}

func (fsm *FSMHandler) onRollback(endpoints []*WeightedEndpoint) error {
	if !fsm.timerExpired() {
		return nil
	}
	fsm.state = FSMStart
	return nil
}

func (fsm *FSMHandler) setTimer() {
	fsm.timer = fsm.timeProvider.UtcNow().Add(fsm.backoffDuration)
}

func (fsm *FSMHandler) timerExpired() bool {
	return fsm.timer.Before(fsm.timeProvider.UtcNow())
}

// Splits endpoint into two groups of endpoints with bad performance and good performance. It does compare relative
// performances of the endpoints though, so if all endpoints have the same performance,
func splitEndpoints(endpoints []*WeightedEndpoint) (good []*WeightedEndpoint, bad []*WeightedEndpoint) {
	avg := avgFailRate(endpoints)
	for _, e := range endpoints {
		if greater(e.failRateMeter.GetRate(), avg) {
			bad = append(bad, e)
		} else {
			good = append(good, e)
		}
	}
	return good, bad
}

func metricsReady(endpoints []*WeightedEndpoint) bool {
	for _, e := range endpoints {
		if !e.failRateMeter.IsReady() {
			return false
		}
	}
	return true
}

func adjustWeights(good, bad []*WeightedEndpoint) []*changedEndpoint {
	changedEndpoints := make([]*changedEndpoint, len(good))
	for i, e := range good {
		changed := &changedEndpoint{
			weightBefore:    e.GetEffectiveWeight(),
			failRatioBefore: e.failRateMeter.GetRate(),
			endpoint:        e,
		}
		changedEndpoints[i] = changed
		if increase(e.GetEffectiveWeight()) < FSMMaxWeight {
			e.setEffectiveWeight(increase(e.GetEffectiveWeight()))
			log.Infof("FSMHandler updated weight %s", e)
		}
	}
	return changedEndpoints
}

// Compare two fail rates by neglecting the insignificant differences
func greater(a, b float64) bool {
	log.Infof("Greater %f %f, %d %d", a, b, math.Floor(a*10), math.Floor(b*10))
	return math.Floor(a*10) > math.Ceil(b*10)
}

func avgFailRate(endpoints []*WeightedEndpoint) float64 {
	r := float64(0)
{'fraction_non_alphanumeric': 0.09864001310830739, 'fraction_numerical': 0.004669834507619204, 'mean_word_length': 3.9063504823151125, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6619807', 'n_tokens_mistral': 4248, 'n_tokens_neox': 3947, 'n_words': 1224}