multinet/balancer.go

62 lines
1.5 KiB
Go
Raw Normal View History

package multinet
import (
"context"
"errors"
"net"
"sync/atomic"
)
// BalancerType reperents the algorithm which is used to pick source address.
type BalancerType string
const (
// BalancerTypeNoop picks first address for which link is up.
BalancerTypeNoop BalancerType = ""
// BalancerTypeNoop implements simple round-robin between up links.
// It is not fair in case some links are down.
BalancerTypeRoundRobin BalancerType = "roundrobin"
)
type balancer interface {
DialContext(ctx context.Context, s *Subnet, network, address string) (net.Conn, error)
}
type roundRobin struct {
d *dialer
i atomic.Uint32
}
func (r *roundRobin) DialContext(ctx context.Context, s *Subnet, network, address string) (net.Conn, error) {
next := int(r.i.Add(1))
for i := range s.Interfaces {
ii := s.Interfaces[(i+next)%len(s.Interfaces)]
if ii.Down {
continue
}
dd := r.d.dialer
dd.LocalAddr = ii.LocalAddr
return r.d.dialContext(&dd, ctx, "tcp", address)
}
return nil, errors.New("(*roundRobin).DialContext: no suitale node found")
}
type firstEnabled struct {
d *dialer
}
func (r *firstEnabled) DialContext(ctx context.Context, s *Subnet, network, address string) (net.Conn, error) {
for i := range s.Interfaces {
ii := s.Interfaces[i%len(s.Interfaces)]
if ii.Down {
continue
}
dd := r.d.dialer
dd.LocalAddr = ii.LocalAddr
return r.d.dialContext(&dd, ctx, "tcp", address)
}
return nil, errors.New("(*firstEnabled).DialContext: no suitale node found")
}