<commit_before>package balancer

import (
	"fmt"
	"math/rand"
	"net"
	"sort"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("balancer")
)

var (
	emptyDialers = []*dialer{}
)

// Balancer balances connections established by one or more Dialers.
type Balancer struct {
	dialers []*dialer
}

// New creates a new Balancer using the supplied Dialers.
func New(dialers ...*Dialer) *Balancer {
	dhs := make([]*dialer, 0, len(dialers))
	for _, d := range dialers {
		dl := &dialer{Dialer: d}
		dl.start()
		dhs = append(dhs, dl)
	}
	// Sort dialers by QOS (ascending) for later selection
	sort.Sort(byQOSAscending(dhs))
	return &Balancer{
		dialers: dhs,
	}
}

// DialQOS dials network, addr using one of the currently active configured
// Dialers. It attempts to use a Dialer whose QOS is higher than targetQOS, but
// will use the highest QOS Dialer(s) if none meet targetQOS. When multiple
// Dialers meet the targetQOS, load is distributed amongst them randomly based
// on their relative Weights.
//
// If a Dialer fails to connect, Dial will keep falling back through the
// remaining Dialers until it either manages to connect, or runs out of dialers
// in which case it returns an error.
func (b *Balancer) DialQOS(network, addr string, targetQOS int) (net.Conn, error) {
	dialers := b.dialers
	for {
		if len(dialers) == 0 {
			return nil, fmt.Errorf("No dialers left to try")
		}
		var d *dialer
		d, dialers = randomDialer(dialers, targetQOS)
		if d == nil {
			return nil, fmt.Errorf("No dialers left")
		}
		if d.Label != "" {
			log.Debugf("Dialing %s://%s with %s", network, addr, d.Label)
		}
		conn, err := d.Dial(network, addr)
		if err != nil {
			log.Tracef("Unable to dial: %s", err)
			d.onError(err)
			continue
		}
		return conn, nil
	}
}

// Dial is like DialQOS with a targetQOS of 0.
func (b *Balancer) Dial(network, addr string) (net.Conn, error) {
	return b.DialQOS(network, addr, 0)
}

// Close closes this Balancer, stopping all background processing. You must call
// Close to avoid leaking goroutines.
func (b *Balancer) Close() {
	for _, d := range b.dialers {
		d.stop()
	}
}

func randomDialer(dialers []*dialer, targetQOS int) (chosen *dialer, others []*dialer) {
	// Weed out inactive dialers and those with too low QOS
	filtered := make([]*dialer, 0)
	highestQOS := 0
	for _, d := range dialers {
		if !d.isActive() {
			log.Trace("Excluding inactive dialer")
			continue
		}

		highestQOS = d.QOS // don't need to compare since dialers are already sorted by QOS (ascending)
		if d.QOS >= targetQOS {
			log.Tracef("Including dialer with QOS %d meeting targetQOS %d", d.QOS, targetQOS)
			filtered = append(filtered, d)
		}
	}

	if len(filtered) == 0 {
		log.Trace("No dialers meet targetQOS, finding remaining with highest QOS")
		for _, d := range dialers {
			if !d.isActive() {
				log.Trace("Excluding inactive dialer")
				continue
			}

			if d.QOS == highestQOS {
				filtered = append(filtered, d)
			}
		}
	}

	if len(filtered) == 0 {
		log.Trace("Still no dialers!")
		return nil, nil
	}

	totalWeights := 0
	for _, d := range filtered {
		totalWeights += d.Weight
	}

	// Pick a random server using a target value between 0 and the total weights
	t := rand.Intn(totalWeights)
	aw := 0
	for _, d := range filtered {
		aw += d.Weight
		if aw > t {
			log.Trace("Reached random target value, using this dialer")
			return d, withoutDialer(dialers, d)
		}
	}

	// We should never reach this
	panic("No dialer found!")
}

func withoutDialer(dialers []*dialer, d *dialer) []*dialer {
	for i, existing := range dialers {
		if existing == d {
			return without(dialers, i)
		}
	}
	log.Tracef("Dialer not found for removal: %s", d)
	return dialers
}

func without(dialers []*dialer, i int) []*dialer {
	if len(dialers) == 1 {
		return emptyDialers
	} else if i == len(dialers)-1 {
		return dialers[:i]
	} else {
		c := make([]*dialer, len(dialers)-1)
		copy(c[:i], dialers[:i])
		copy(c[i:], dialers[i+1:])
		return c
	}<commit_msg>Made logic for picking random dialer more readable<commit_after>package balancer

import (
	"fmt"
	"math/rand"
	"net"
	"sort"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("balancer")
)

var (
	emptyDialers = []*dialer{}
)

// Balancer balances connections established by one or more Dialers.
type Balancer struct {
	dialers []*dialer
}

// New creates a new Balancer using the supplied Dialers.
func New(dialers ...*Dialer) *Balancer {
	dhs := make([]*dialer, 0, len(dialers))
	for _, d := range dialers {
		dl := &dialer{Dialer: d}
		dl.start()
		dhs = append(dhs, dl)
	}
	// Sort dialers by QOS (ascending) for later selection
	sort.Sort(byQOSAscending(dhs))
	return &Balancer{
		dialers: dhs,
	}
}

// DialQOS dials network, addr using one of the currently active configured
// Dialers. It attempts to use a Dialer whose QOS is higher than targetQOS, but
// will use the highest QOS Dialer(s) if none meet targetQOS. When multiple
// Dialers meet the targetQOS, load is distributed amongst them randomly based
// on their relative Weights.
//
// If a Dialer fails to connect, Dial will keep falling back through the
// remaining Dialers until it either manages to connect, or runs out of dialers
// in which case it returns an error.
func (b *Balancer) DialQOS(network, addr string, targetQOS int) (net.Conn, error) {
	dialers := b.dialers
	for {
		if len(dialers) == 0 {
			return nil, fmt.Errorf("No dialers left to try")
		}
		var d *dialer
		d, dialers = randomDialer(dialers, targetQOS)
		if d == nil {
			return nil, fmt.Errorf("No dialers left")
		}
		if d.Label != "" {
			log.Debugf("Dialing %s://%s with %s", network, addr, d.Label)
		}
		conn, err := d.Dial(network, addr)
		if err != nil {
			log.Tracef("Unable to dial: %s", err)
			d.onError(err)
			continue
		}
		return conn, nil
	}
}

// Dial is like DialQOS with a targetQOS of 0.
func (b *Balancer) Dial(network, addr string) (net.Conn, error) {
	return b.DialQOS(network, addr, 0)
}

// Close closes this Balancer, stopping all background processing. You must call
// Close to avoid leaking goroutines.
func (b *Balancer) Close() {
	for _, d := range b.dialers {
		d.stop()
	}
}

func randomDialer(dialers []*dialer, targetQOS int) (chosen *dialer, others []*dialer) {
	// Weed out inactive dialers and those with too low QOS
	filtered, highestQOS := dialersMeetingQOS(dialers, targetQOS)

	if len(filtered) == 0 {
		log.Trace("No dialers meet targetQOS, using those with highest QOS")
		filtered, _ = dialersMeetingQOS(dialers, highestQOS)
	}

	if len(filtered) == 0 {
		log.Trace("Still no dialers!")
		return nil, nil
	}

	totalWeights := 0
	for _, d := range filtered {
		totalWeights += d.Weight
	}

	// Pick a random server using a target value between 0 and the total weights
	t := rand.Intn(totalWeights)
	aw := 0
	for _, d := range filtered {
		aw += d.Weight
		if aw > t {
			log.Trace("Reached random target value, using this dialer")
			return d, withoutDialer(dialers, d)
		}
	}

	// We should never reach this
	panic("No dialer found!")
}

func dialersMeetingQOS(dialers []*dialer, targetQOS int) ([]*dialer, int) {
	filtered := make([]*dialer, 0)
	highestQOS := 0
	for _, d := range dialers {
		if !d.isActive() {
			log.Trace("Excluding inactive dialer")
			continue
		}

		highestQOS = d.QOS // don't need to compare since dialers are already sorted by QOS (ascending)
		if d.QOS >= targetQOS {
			log.Tracef("Including dialer with QOS %d meeting targetQOS %d", d.QOS, targetQOS)
			filtered = append(filtered, d)
		}
	}

	return filtered, highestQOS
}

func withoutDialer(dialers []*dialer, d *dialer) []*dialer {
	for i, existing := range dialers {
		if existing == d {
			return without(dialers, i)
		}
	}
	log.Tracef("Dialer not found for removal: %s", d)
	return dialers
}

func without(dialers []*dialer, i int) []*dialer {
	if len(dialers) == 1 {
		return emptyDialers
	} else if i == len(dialers)-1 {
		return dialers[:i]
	} else {
		c := make([]*dialer, len(dialers)-1)
		copy(c[:i], dialers[:i])
		copy(c[i:], dialers[i+1:])
		return c
	}
{'fraction_non_alphanumeric': 0.1286993285252425, 'fraction_numerical': 0.0037304153195722458, 'mean_word_length': 3.5134680134680134, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '5273054', 'n_tokens_mistral': 3203, 'n_tokens_neox': 2895, 'n_words': 952}