mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-09-01 06:14:01 +03:00
Add load balancing network dialer
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
package network
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math/rand"
|
||||
"net"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
type loadBalancedDialer struct {
|
||||
dialers []Dialer
|
||||
}
|
||||
|
||||
func (l loadBalancedDialer) Dial(network, address string) (net.Conn, error) {
|
||||
return l.DialContext(context.Background(), network, address)
|
||||
}
|
||||
|
||||
func (l loadBalancedDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
length := len(l.dialers)
|
||||
start := rand.Intn(length)
|
||||
moved := false
|
||||
|
||||
for i := start; i != start || !moved; i = (i + 1) % length {
|
||||
moved = true
|
||||
if conn, err := l.dialers[i].DialContext(ctx, network, address); err == nil {
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, ErrCannotDialWithAllProxies
|
||||
}
|
||||
|
||||
func NewLoadBalancedDialer(baseDialer Dialer, proxyURLs []*url.URL) (Dialer, error) {
|
||||
switch len(proxyURLs) {
|
||||
case 0:
|
||||
return baseDialer, nil
|
||||
case 1:
|
||||
return NewSocks5Dialer(baseDialer, proxyURLs[0])
|
||||
}
|
||||
|
||||
dialers := []Dialer{}
|
||||
|
||||
for _, u := range proxyURLs {
|
||||
dialer, err := NewSocks5Dialer(newProxyDialer(baseDialer, u), u)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
dialers = append(dialers, dialer)
|
||||
}
|
||||
|
||||
return loadBalancedDialer{
|
||||
dialers: dialers,
|
||||
}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user