56 lines
1.5 KiB
Go
56 lines
1.5 KiB
Go
package multinet
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"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"
|
|
)
|
|
|
|
var errNoSuitableNodeFound = errors.New("no suitale node found")
|
|
|
|
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.SourceIPs {
|
|
ii := s.SourceIPs[(i+next)%len(s.SourceIPs)]
|
|
dd := r.d.dialer
|
|
dd.LocalAddr = &net.TCPAddr{IP: net.IP(ii.AsSlice())}
|
|
return r.d.dialContext(ctx, &dd, network, address)
|
|
}
|
|
return nil, fmt.Errorf("(*roundRobin).DialContext: %w", errNoSuitableNodeFound)
|
|
}
|
|
|
|
type firstEnabled struct {
|
|
d *dialer
|
|
}
|
|
|
|
func (r *firstEnabled) DialContext(ctx context.Context, s *Subnet, network, address string) (net.Conn, error) {
|
|
for i := range s.SourceIPs {
|
|
ii := s.SourceIPs[i]
|
|
dd := r.d.dialer
|
|
dd.LocalAddr = &net.TCPAddr{IP: net.IP(ii.AsSlice())}
|
|
return r.d.dialContext(ctx, &dd, network, address)
|
|
}
|
|
return nil, fmt.Errorf("(*firstEnabled).DialContext: %w", errNoSuitableNodeFound)
|
|
}
|