multinet/balancer.go

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")
}